Commits

Branch selector

 main

User selector

yashk2810

Datepicker

All time

Commit History

Commits on Dec 18, 2024

• Remove jax_enable_memories flag

  Show description for f7bece0
  

  yashk2810

  authored and

  Orbax Authors

  committed
  f7bece0

  Copy full SHA for f7bece0

Commits on Oct 7, 2024

• Now that JAX supports donate on device_put, use that instead of the ugly device_put -> block -> delete routine

  Show description for db46fc7
  

  yashk2810

  authored and

  Orbax Authors

  committed
  db46fc7

  Copy full SHA for db46fc7

Commits on Aug 7, 2024

• Add memory_kind to the PositionalShardingMetadata in test_convert_between_jax_positional_sharding_and_sharding_metadata test

  Show description for a301fc8
  

  yashk2810

  authored and

  Orbax Authors

  committed
  a301fc8

  Copy full SHA for a301fc8

Commits on Jul 15, 2023

• Internal change

  Show description for 4e3a99f
  

  yashk2810

  authored and

  copybara-github

  committed
  4e3a99f

  Copy full SHA for 4e3a99f

Commits on May 12, 2023

• Internal change

  Show description for df7810f
  

  yashk2810

  authored and

  copybara-github

  committed
  df7810f

  Copy full SHA for df7810f

Commits on Mar 10, 2023

• Internal change

  Show description for e1b15c1
  

  yashk2810

  authored and

  copybara-github

  committed
  e1b15c1

  Copy full SHA for e1b15c1

Commits on Nov 17, 2022

• Internal change

  Show description for 5459ae8
  

  yashk2810

  authored and

  copybara-github

  committed
  5459ae8

  Copy full SHA for 5459ae8

Commits on Oct 31, 2022

• Support jax.Array in T5X. This feature is currently disabled but can be enabled by setting use_gda=True and use_jax_array=True.

  Show description for a91af40
  

  yashk2810

  authored and

  copybara-github

  committed
  a91af40

  Copy full SHA for a91af40

Commits on Oct 18, 2022

• If both jax.Array and GDA are enabled, disable GDA and give priority to Array.

  Show description for d230f3a
  

  yashk2810

  authored and

  copybara-github

  committed
  d230f3a

  Copy full SHA for d230f3a

Commits on Oct 14, 2022

• Internal change

  Show description for 3b1055e
  

  yashk2810

  authored and

  copybara-github

  committed
  3b1055e

  Copy full SHA for 3b1055e

Commits on Oct 10, 2022

• Disable GDA when jax.Array is enabled and use addressable_data since local_data is deprecated.

  Show description for 809055a
  

  yashk2810

  authored and

  copybara-github

  committed
  809055a

  Copy full SHA for 809055a

Commits on Oct 6, 2022

• Take shardings as a parameter to deserialize and run_deserialization instead of mesh and pspecs.

  Show description for be9fe35
  

  yashk2810

  authored and

  copybara-github

  committed
  be9fe35

  Copy full SHA for be9fe35

Commits on Sep 27, 2022

• Move array.py and sharding.py from experimental/ to _src/.

  Show description for ce201e5
  

  yashk2810

  authored and

  copybara-github

  committed
  ce201e5

  Copy full SHA for ce201e5

Commits on Sep 26, 2022

• Use public jax.Array type rather than the internal type.

  Show description for 32befc9
  

  yashk2810

  authored and

  copybara-github

  committed
  32befc9

  Copy full SHA for 32befc9