Commits

Branch selector

 main

User selector

yashk2810

Datepicker

All time

Commit History

Commits on Feb 16, 2024

• Add a benchmark measuring device_put's speed for a 4GB input array

  Show description for 0b542ff
  

  yashk2810

  authored and

  jax authors

  committed
  0b542ff

  Copy full SHA for 0b542ff

Commits on Feb 15, 2024

• Update the cuda 12 dependencies since we upgraded to cuda 12.3

  Show description for f125509
  

  yashk2810

  authored and

  jax authors

  committed
  f125509

  Copy full SHA for f125509

• Partial rollback xla translation APIs that were removed in 0.4.24 release

  Show description for 8888006
  

  yashk2810

  authored and

  jax authors

  committed
  8888006

  Copy full SHA for 8888006

Commits on Feb 14, 2024

• Remove the access of `mesh.devices` from shard_map in tracing mode. It's okay to access it in eager mode

  Show description for da32609
  

  yashk2810

  authored and

  jax authors

  committed
  da32609

  Copy full SHA for da32609

Commits on Feb 9, 2024

• Add `offload_dot_with_no_batch_dims` to `jax.checkpoint_policies`

  Show description for 6e17bb3
  

  yashk2810

  authored and

  jax authors

  committed
  6e17bb3

  Copy full SHA for 6e17bb3

• Rename `tpu_hbm` to `device` so that GPU can also use `device` which will allow users to have device portable code.

  Show description for 486436f
  

  yashk2810

  authored and

  jax authors

  committed
  486436f

  Copy full SHA for 486436f

• Update changelog and install doc to mention cuda 12.3 switch

  Show description for 73e3ded
  

  yashk2810

  authored and

  jax authors

  committed
  73e3ded

  Copy full SHA for 73e3ded

Commits on Feb 6, 2024

• Skip activation offload tests on GPU and CPU until they are supported

  Show description for 299b983
  

  yashk2810

  authored and

  jax authors

  committed
  299b983

  Copy full SHA for 299b983

• Support activation offloading to host in JAX!

  Show description for d30c3ad
  

  yashk2810

  authored and

  jax authors

  committed
  d30c3ad

  Copy full SHA for d30c3ad

Commits on Jan 30, 2024

• Add sharding to ShapeDtypeStruct retured by eval_shape if `jit` has `out_shardings` specified

  Show description for d9122b8
  

  yashk2810

  authored and

  jax authors

  committed
  d9122b8

  Copy full SHA for d9122b8

Commits on Jan 24, 2024

• Add an internal `_device_list` parameter to GSPMDSharding so that we can save on the initialization cost of PyDeviceList when creating `GSPMDSharding` from other shardings

  Show description for a63197f
  

  yashk2810

  authored and

  jax authors

  committed
  a63197f

  Copy full SHA for a63197f

• Preserve single device NamedSharding/PositionalSharding on the output instead of always return SingleDeviceShardings.

  Show description for 6f96c96
  

  yashk2810

  authored and

  jax authors

  committed
  6f96c96

  Copy full SHA for 6f96c96

Commits on Jan 20, 2024

• Internal change

  Show description for 3a0b495
  

  yashk2810

  authored and

  jax authors

  committed
  3a0b495

  Copy full SHA for 3a0b495

Commits on Jan 19, 2024

• Make eval_shape a wrapper around `jax.jit(f).eval_shape(*args, **kwargs)`

  Show description for f04f305
  

  yashk2810

  authored and

  jax authors

  committed
  f04f305

  Copy full SHA for f04f305

Commits on Jan 18, 2024

• Use jit's jaxpr creation function for eval_shape to maximize tracing cache hits.

  Show description for 51ef738
  

  yashk2810

  authored and

  jax authors

  committed
  51ef738

  Copy full SHA for 51ef738

• Replace _DeviceAssignment with xc.DeviceList

  Show description for 71c9be1
  

  yashk2810

  authored and

  jax authors

  committed
  71c9be1

  Copy full SHA for 71c9be1

• Add `__str__` to `Mesh` so that in jaxprs the mesh doesn't print all the device ids.

  Show description for 72f00eb
  

  yashk2810

  authored and

  jax authors

  committed
  72f00eb

  Copy full SHA for 72f00eb

Commits on Jan 8, 2024

• Preserve the `spec`s passed by the user in the output sharding from a eager shard_map.

  Show description for dccc0e8
  

  yashk2810

  authored and

  jax authors

  committed
  dccc0e8

  Copy full SHA for dccc0e8

Commits on Jan 5, 2024

• Remove indices and devices from shard_arg_handlers and shard_args.

  Show description for b8098b1
  

  yashk2810

  authored and

  jax authors

  committed
  b8098b1

  Copy full SHA for b8098b1

• Remove the check for existence of _npy_value before taking the fast path for `__getitem__`. This must have been a remnant of SDA era.

  Show description for 085b23d
  

  yashk2810

  authored and

  jax authors

  committed
  085b23d

  Copy full SHA for 085b23d

Commits on Jan 4, 2024

• Rename batch_dims variable in shard_alike

  Show description for 7bba80d
  

  yashk2810

  authored and

  jax authors

  committed
  7bba80d

  Copy full SHA for 7bba80d

• Add correct batching rule for shard_alike

  Show description for 68abe0d
  

  yashk2810

  authored and

  jax authors

  committed
  68abe0d

  Copy full SHA for 68abe0d

Commits on Jan 2, 2024

• Delete sharding spec to HloSharding conversion since it's not used anymore.

  Show description for c0d4653
  

  yashk2810

  authored and

  jax authors

  committed
  c0d4653

  Copy full SHA for c0d4653

• Delete `mesh_sharding_specs` from JAX

  Show description for 5192fca
  

  yashk2810

  authored and

  jax authors

  committed
  5192fca

  Copy full SHA for 5192fca

• Remove reliance on `ShardingSpec`s from `NamedSharding` to `HloSharding` conversion.

  Show description for 697f17a
  

  yashk2810

  authored and

  jax authors

  committed
  697f17a

  Copy full SHA for 697f17a

Commits on Dec 29, 2023

• Add shape of PositionalSharding to it's repr

  Show description for 4b76d03
  

  yashk2810

  authored and

  jax authors

  committed
  4b76d03

  Copy full SHA for 4b76d03

Commits on Dec 26, 2023

• Add shard_alike tf_impl rule

  Show description for 9c2b21f
  

  yashk2810

  authored and

  jax authors

  committed
  9c2b21f

  Copy full SHA for 9c2b21f

Commits on Dec 24, 2023

• Replace usage of make_array_from_callback in full_like with shard_alike

  Show description for 0d6007c
  

  yashk2810

  authored and

  jax authors

  committed
  0d6007c

  Copy full SHA for 0d6007c

Commits on Dec 22, 2023

• Cache ClosedJaxpr creation to minimize cache_misses. ClosedJaxpr should always be created under a cache.

  Show description for 4c9241e
  

  yashk2810

  authored and

  jax authors

  committed
  4c9241e

  Copy full SHA for 4c9241e

Commits on Dec 21, 2023

• Cleanup _find_arg_mismatch logic

  Show description for 9792e00
  

  yashk2810

  authored and

  jax authors

  committed
  9792e00

  Copy full SHA for 9792e00

• Always return a NamedSharding from eager shard_map

  Show description for 57d74d6
  

  yashk2810

  authored and

  jax authors

  committed
  57d74d6

  Copy full SHA for 57d74d6

Commits on Dec 20, 2023

• Always flatten args and kwargs together i.e. `tree_flatten((args, kwargs))` so that we have a uniform in_tree structure everywhere.

  Show description for 90e47fb
  

  yashk2810

  authored and

  jax authors

  committed
  90e47fb

  Copy full SHA for 90e47fb

• Expose `shard_alike` via jax.experimental. The API is `x, y = shard_like(x, y)`.

  Show description for 72fbdb2
  

  yashk2810

  authored and

  jax authors

  committed
  72fbdb2

  Copy full SHA for 72fbdb2

Commits on Dec 19, 2023

• Call shard_arg fallback in pjit's cpp fast path instead of dropping out completely.

  Show description for 9b6bf2c
  

  yashk2810

  authored and

  jax authors

  committed
  9b6bf2c

  Copy full SHA for 9b6bf2c

Commits on Dec 14, 2023

• [roll forward 2] Remove the `jax_require_devices_during_lowering flag since it was temporary. Added the semi-breaking change to Changelog.md.

  Show description for 8bf3a86
  

  yashk2810

  authored and

  jax authors

  committed
  8bf3a86

  Copy full SHA for 8bf3a86

Pagination

Previous

Next