Commits

Branch selector

 main

User selector

yashk2810

Datepicker

All time

Commit History

Commits on Jul 26, 2025

• Don't use aval's mesh if any of the axes in it as manual

  Show description for 413c976
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  413c976

  Copy full SHA for 413c976

Commits on Jul 25, 2025

• Fix `_add_transpose` rule to make sure that cotangent sharding matches the primal shardings.

  Show description for 2202f95
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  2202f95

  Copy full SHA for 2202f95

• Remove `eval_shape` as a method on Traced and rather use `.out_info` because `.trace` already evals.

  Show description for 404981f
  

  2 peopleauthored and

  Google-ML-Automation

  committed
  404981f

  Copy full SHA for 404981f

Commits on Jul 24, 2025

• [JAX] Move 7x mesh creation code to OSS.

  Show description for e62b4ab
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  e62b4ab

  Copy full SHA for e62b4ab

• Make `tiling` argument of `Layout` class public and the property too. So basically it's `s/_tiling/tiling/g`.

  Show description for cbecf85
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  cbecf85

  Copy full SHA for cbecf85

Commits on Jul 23, 2025

• Cache `canonicalize_sharding` and remove the canonicalization in `_scatter_update` since the out_sharding passed to scatter_update should already be

  Show description for 768a64e
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  768a64e

  Copy full SHA for 768a64e

• Remove out_sharding from scatter_add because we now have a scatter sharding rule

  Show description for 1c0f24b
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  1c0f24b

  Copy full SHA for 1c0f24b

• Fix a bug where doing a CPU jax.Array -> TPU jax.Array transfer was not setting the layout correctly on the output TPU array.

  Show description for b57f077
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  b57f077

  Copy full SHA for b57f077

Commits on Jul 22, 2025

• Remove JAX_IFRT_VERSION_NUMBER guards after the release of 0.7.0

  Show description for fb1a8fc
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  fb1a8fc

  Copy full SHA for fb1a8fc

• #sdy Delete JAX test configs enabling Shardy.

  Show description for ff595f6
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  ff595f6

  Copy full SHA for ff595f6

• Remove `vmap` and `pmap` mentions in the rst section title.

  Show description for 430cc85
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  430cc85

  Copy full SHA for 430cc85

• Add `smap` to the API docs on the website. Fixes https://github.com/jax-ml/jax/issues/30407

  Show description for 5be0ee7
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  5be0ee7

  Copy full SHA for 5be0ee7

• Remove jaxlib_extension_version, ifrt_version guards now that 0.7.0 release is out

  Show description for 5a6dd76
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  5a6dd76

  Copy full SHA for 5a6dd76

• Remove `SafeNumDevices` now that all shardings are implemented in C++

  Show description for 3eaa7f9
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  3eaa7f9

  Copy full SHA for 3eaa7f9

Commits on Jul 21, 2025

• Fix mesh_cast + new rng keys usage

  Show description for 5371d6b
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  5371d6b

  Copy full SHA for 5371d6b

• If the input aval mesh is empty, use the mesh of out_sharding passed to with_sharding_constraint on the output aval. We only use the mesh and not the pspec. This leads to more tracing cache hits as…

  Show description for bd32431
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  bd32431

  Copy full SHA for bd32431

• Consider this example:

  Show description for b97ad42
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  b97ad42

  Copy full SHA for b97ad42

Commits on Jul 19, 2025

• Separate out `vma` and `unreduced/reduced` in the ShapedArray str_short().

  Show description for 0acbd3b
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  0acbd3b

  Copy full SHA for 0acbd3b

• Use `operand.aval.sharding` in convert_element_type's transpose rule

  Show description for dd36a9f
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  dd36a9f

  Copy full SHA for dd36a9f

• Remove the assert from scatter_addsub_transpose_rule. Fixes https://github.com/jax-ml/jax/issues/30327. We don't need to use the `out_sharding` in transpose rule anyways that's bound to the primitive.

  Show description for 59e15b9
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  59e15b9

  Copy full SHA for 59e15b9

Commits on Jul 18, 2025

• Move ShapeDtypeStruct to core.py to break circular deps

  Show description for fec027e
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  fec027e

  Copy full SHA for fec027e

• Replace `jax.stages.OutInfo` with `jax.ShapeDtypeStruct`

  Show description for eb194e9
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  eb194e9

  Copy full SHA for eb194e9

• Add `eval_shape` as a method of `Traced`.

  Show description for dd59b47
  

  2 peopleauthored and

  Google-ML-Automation

  committed
  dd59b47

  Copy full SHA for dd59b47

• If you are doing `jax.eval_shape` on a jitted function, call `jitted_fn.eval_shape(*args)` internally as a shortcut.

  Show description for 74fdb41
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  74fdb41

  Copy full SHA for 74fdb41

Commits on Jul 16, 2025

• Set `check_leaks=False` in `def direct_linearize` to fix some tests

  Show description for 4752cc9
  

  3 peopleauthored and

  Google-ML-Automation

  committed
  4752cc9

  Copy full SHA for 4752cc9

• fix jax2tf scatter impl

  Show description for 45f67c7
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  45f67c7

  Copy full SHA for 45f67c7

Commits on Jul 15, 2025

• Add a test mentioned in https://github.com/jax-ml/jax/issues/29883 to JAX test suite

  Show description for 9b07a16
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  9b07a16

  Copy full SHA for 9b07a16

Commits on Jul 14, 2025

• If operand and indices are fully replicated, allow that case to work in gather.

  Show description for 88b82cf
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  88b82cf

  Copy full SHA for 88b82cf

• Add `out_sharding` to `scatter_add_p` so that we can pass a value to that from gather's transpose rule thereby supporting sharding_in_types in gather transpose.

  Show description for 3f380b7
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  3f380b7

  Copy full SHA for 3f380b7

• Error out if compiler_options are specified in a nested jit. They can only be specified on the top level jit.

  Show description for 6f8456e
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  6f8456e

  Copy full SHA for 6f8456e

Commits on Jul 12, 2025

• Enable decorator factory pattern for `auto_axes` and `explicit_axes` API. Now users can do: `@auto_axes(...)` and `@explicit_axes(...)`

  Show description for 440f31f
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  440f31f

  Copy full SHA for 440f31f

Commits on Jul 11, 2025

• Allow dropping into explicit_axes with a reshard in an outer vmap with spmd_axis_name set.

  Show description for f926b3c
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  f926b3c

  Copy full SHA for f926b3c

• Fix pjit failures with direct linearize

  Show description for 4433c14
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  4433c14

  Copy full SHA for 4433c14

• Fixes `linalg.inv` + shard_map with check_vma=True by adding `standard_insert_pvary` in `linalg.solve`. Fixes https://github.com/jax-ml/jax/issues/30157

  Show description for a06e9e3
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  a06e9e3

  Copy full SHA for a06e9e3

• The repr of Literal should mention the class name "Literal" in it

  Show description for 6b49571
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  6b49571

  Copy full SHA for 6b49571

Pagination

Previous

Next