All pull requests

Pull requests

• Fix Restore Checkpoint loading with partial weight By setting restore_args to None if the corresponding state value is None.

  #1639·copybara-service[bot] opened Jan 15, 2026Bot·

  ·
  1

• Fix a bug in Orbax checkpointing where None values PyTree are not handled correctly. This is caused by a recent update in jax where None values are no longer considered as a leaf node: https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-34-october-4-2024.

  #1604·copybara-service[bot] opened Oct 31, 2024Bot·

  ·
  1

• Disable golden tests in contributed gpu layers_test.py

  #1600·levskaya opened Oct 21, 2024Collaborator·

  ·

• add wait_until_finished under checkpoint_manager to fix issue when checkpoint was not completed.

  #1581·copybara-service[bot] opened Jul 17, 2024Bot·

  ·

• Sort devices by their implicit order instead of explicitly by id. IDs may be randomly generated, so it's better to rely on the implicit order, which is currently based on (process index, id).

  #1578·copybara-service[bot] opened Jul 10, 2024Bot·

  ·
  1

• Revert get_local_data to original implementation.

  #1571·copybara-service[bot] opened Jun 12, 2024Bot·

  ·

• Fix bug in T5X CheckpointManager argument propagation when SaveBestCheckpoint is set through gin.

  #1568·copybara-service[bot] opened Jun 7, 2024Bot·

  ·
  1

• Propagate metric_name_to_monitor to Orbax CheckpointManager.

  #1566·copybara-service[bot] opened Jun 3, 2024Bot·

  ·

• Added MAXO-GEN models, namely, SAC, AGIQA-1k, AGIQA-3k and AIGCIQA-2023.

  #1543·copybara-service[bot] opened Apr 19, 2024Bot·

  ·
  1

• Added MAXO-GEN models, namely, SAC, AGIQA-1k, AGIQA-3k and AIGCIQA-2023.

  #1542·copybara-service[bot] opened Apr 19, 2024Bot·

  ·
  1

• jax.tree.map->jax.tree_map

  #1535·jntdst opened Mar 28, 2024·

  ·
  1

• DNS test cloud issue.

  #1531·copybara-service[bot] opened Mar 20, 2024Bot·

  ·

• internal

  #1469·copybara-service[bot] opened Dec 17, 2023Bot·

  ·
  1

• internal

  #1468·copybara-service[bot] opened Dec 17, 2023Bot·

  ·
  1

• internal

  #1467·copybara-service[bot] opened Dec 17, 2023Bot·

  ·
  1

• internal

  #1466·copybara-service[bot] opened Dec 17, 2023Bot·

  ·
  1

• Remove tensorflow-cpu dependency

  #1464·vedansh0739 opened Dec 15, 2023·

  ·
  1

• Replace references to deprecated device_buffer attributes

  #1447·copybara-service[bot] opened Dec 6, 2023Bot·

  ·

• Unpins tfds-nightly in gpu extra which is in conflict with seqio's specification

  #1416·terrykong opened Oct 11, 2023Contributor·

  ·

• Add a jax context manager to prevent crashes in some cases.

  #1402·copybara-service[bot] opened Sep 20, 2023Bot·

  ·
  1

• Add DALI support to t5x

  #1393·terrykong opened Sep 8, 2023Contributor·

  ·
  1

• Add support for partial checkpoint restore

  #1392·terrykong opened Sep 8, 2023Contributor·

  ·
  1

• [NVIDIA] Added transformer engine support and GPU optimizations

  #1391·terrykong opened Sep 8, 2023Contributor·

  ·
  1

• Fix broken type annotations, and replace pytype disables with typing.cast

  #1384·copybara-service[bot] opened Aug 25, 2023Bot·

  ·
  1

• Typo fixes

  #1381·copybara-service[bot] opened Aug 23, 2023Bot·

  ·
  1