jax

CI: 03/14/25 upstream sync

#285

Merged

Support nvfp4

Update gammainc and gammaincc against scipy 1.16: return nan whenever one of operands is nan.

Merge pull request #27041 from carlosgmartin:fix_binomial_value_error

Merge pull request #27092 from pearu:pearu/gammainc-bug-fix

[mosaic_gpu/pallas:mgpu] Eradicate wgmma_layout

[Pallas] Plumb compiler flags through source mapper.

[sharding_in_types] Make `reshard` work with np.array.

Add navigation breadcrumbs to docs.

Add a set_mesh API to `jax.sharding`. `set_mesh` sets the sharding and never unsets it i.e. this is just `__enter__` of a ctx manager without `__exit__`

[Mosaic TPU][NFC] Throw NYI error instead of crash when squeeze ref to 1d.

Merge pull request #27044 from carlosgmartin:add_breadcrumbs_to_docs

Add jax.random.multinomial.

[sharding_in_types] Rework the `axis_types` argument in Mesh and AbstractMesh APIs. The changes are:

Make the signature of AbstractMesh to be `AbstractMesh(axis_size: tuple[int, ...], axis_name: tuple[str, ...], *, axis_types)` instead of `AbstractMesh(shape_tuple: tuple[tuple[str, int], ...], *, axis_types)` so that we are consistent across all Mesh APIs: `Mesh`, `AbstractMesh` and `make_mesh`

Rename some internal APIs (set_abstract_mesh -> use_abstract_mesh and set_concrete_mesh -> use_concrete_mesh)

Update XLA dependency to use revision

Add custom job names to group different matrix combinations in the Actions dashboard

Allow late binding `out_shardings` and `in_shardings` in `auto_axes` and `explicit_axes` API

[JAX] Change jax.core.Trace subclasses to call super().__init__().

Merge pull request #26524 from carlosgmartin:random_multinomial

[Mosaic] Improve error messages.

Merge pull request #26944 from wenscarl:wenscarl/nvfp4

[pallas][Mosaic][Easy] Add batch dot dim test, remove check

Raise an error if `jax.config.update('jax_num_cpu_devices', val)` is called after backend is initialized

Remove obsolete fallback for cost analysis.

Reverts 74b4d868e3751c1b4efa315ff8cf771faeb0b663

Precompute the __hash__ of AbstractMesh.

Raise a better error with more info when we see duplicate axis in a PartitionSpec resulting from a sharding rule.

[JAX] Clean up make_array_from_callback_* API benchmarks and add a partially replicated sharding variant

fix shard_map manual mesh axis names with vmap spmd_axis_name

[Mosaic] Allow part of x2 int casts.

`get_sharding` doesn't need to be conditioned on the context mesh

Plumb layout through the creation of IFRT Arrays (roll-forward with fix).

[Mosaic GPU] Introduce an initial transform inference pass.

[Mosaic GPU] Add transform inference rules for `mgpu.async_{load,store}`.

Add explicit support for float8_e4m3b11fnuz in pl.dot

[JAX] Add missing preset for X9 dot optimization on BF16/BF16 -> F32.

[Mosaic GPU][NFC] Simplify implementation for `in_{layout,transforms}_for_operand` utils.

Update XLA dependency to use revision

Precompute a weakref to a Trace≥

Add a variant of safe_map() that has no return value, named foreach().

Remove `internal_test_util` folder and packages from `jax` wheel.

Use "x is y" rather than "id(x) == id(y)".

Install test dependencies from test-requirements.txt instead of requirements.in

[XLA:GPU] Add missing BF16_BF16_F32_X9 matmul option in config.py

[dlpack] Support more DLPack dtypes now that we target DLPack 1.1

Allow `make_array_from_callback` to construct nonaddressable arrays.

Rename AxisTypes enum to AxisType

Support convolution in roofline.

[Pallas] Add legacy PRNG key support to Pallas PRNG

Explicit sharding docs

Merge pull request #27165 from jax-ml:sharding-in-types-doc

Expose `get_abstract_mesh` via the `jax.sharding` namespace

Add `experimental/serialize_executable.py` to `BUILD`.

Fix notebook formatting

[Pallas] Increase #rows when casting to x2.

Merge pull request #27173 from jakevdp:fix-ipynb

[direct-linearize] shmap fixes

[pallas:mosaic_gpu] WG lowering now supports `while_p`

Merge pull request #27015 from mattjj:direct-linearize-fixes-4

Merge remote-tracking branch 'origin/rocm-main' into ci-upstream-sync-147_1

Trivial change for CI