CI: 03/14/25 upstream sync #285
github-actions
merged
63 commits into
rocm-main
from
ci-upstream-sync-147_1
Fix the ValueError message for random.binomial (forgot to use string …
8b6ca564
Support nvfp4
f9aef8a1
Update gammainc and gammaincc against scipy 1.16: return nan whenever…
f608a8c5
Merge pull request #27041 from carlosgmartin:fix_binomial_value_error
e7d10a23
Merge pull request #27092 from pearu:pearu/gammainc-bug-fix
3de7ecf6
[mosaic_gpu/pallas:mgpu] Eradicate wgmma_layout
b34f56bf
[Pallas] Plumb compiler flags through source mapper.
6978f352
[sharding_in_types] Make `reshard` work with np.array.
8674495f
Add navigation breadcrumbs to docs.
bc43b00d
Add a set_mesh API to `jax.sharding`. `set_mesh` sets the sharding an…
47480b44
[Mosaic TPU][NFC] Throw NYI error instead of crash when squeeze ref t…
12c0987e
Merge pull request #27044 from carlosgmartin:add_breadcrumbs_to_docs
ba367cde
Add jax.random.multinomial.
6b69a136
[sharding_in_types] Rework the `axis_types` argument in Mesh and Abst…
c6dcbb67
Make the signature of AbstractMesh to be `AbstractMesh(axis_size: tup…
a4ca0dbc
Rename some internal APIs (set_abstract_mesh -> use_abstract_mesh and…
2d01226b
Update XLA dependency to use revision
c07d839e
Add custom job names to group different matrix combinations in the Ac…
12760af2
Allow late binding `out_shardings` and `in_shardings` in `auto_axes` …
14b9f485
[JAX] Change jax.core.Trace subclasses to call super().__init__().
8effa197
Merge pull request #26524 from carlosgmartin:random_multinomial
bf829ff6
[Mosaic] Improve error messages.
a0f1be12
Merge pull request #26944 from wenscarl:wenscarl/nvfp4
726f49cb
[pallas][Mosaic][Easy] Add batch dot dim test, remove check
47bf22e3
Raise an error if `jax.config.update('jax_num_cpu_devices', val)` is …
e1b62ced
Remove obsolete fallback for cost analysis.
acd6c40f
Reverts 74b4d868e3751c1b4efa315ff8cf771faeb0b663
538a2be7
Precompute the __hash__ of AbstractMesh.
15077544
Raise a better error with more info when we see duplicate axis in a P…
e615e2ac
[JAX] Clean up make_array_from_callback_* API benchmarks and add a pa…
73b8f6ae
fix shard_map manual mesh axis names with vmap spmd_axis_name
34d6bb2e
[Mosaic] Allow part of x2 int casts.
e235fb97
`get_sharding` doesn't need to be conditioned on the context mesh
d3a41d84
Plumb layout through the creation of IFRT Arrays (roll-forward with f…
d7947210
[Mosaic GPU] Introduce an initial transform inference pass.
d028354a
[Mosaic GPU] Add transform inference rules for `mgpu.async_{load,stor…
d09df7c8
Add explicit support for float8_e4m3b11fnuz in pl.dot
cbece0b0
[JAX] Add missing preset for X9 dot optimization on BF16/BF16 -> F32.
43b78c53
[Mosaic GPU][NFC] Simplify implementation for `in_{layout,transforms}…
5098d2ef
Update XLA dependency to use revision
92c57a51
Precompute a weakref to a Trace≥
074216e0
Add a variant of safe_map() that has no return value, named foreach().
8ab33669
Remove `internal_test_util` folder and packages from `jax` wheel.
8fbe3b13
Use "x is y" rather than "id(x) == id(y)".
6fa98fc0
Install test dependencies from test-requirements.txt instead of requi…
5944c9ed
[XLA:GPU] Add missing BF16_BF16_F32_X9 matmul option in config.py
c9ac82c8
[dlpack] Support more DLPack dtypes now that we target DLPack 1.1
97bbc37e
Allow `make_array_from_callback` to construct nonaddressable arrays.
bdb6d033
Rename AxisTypes enum to AxisType
88d4bc3d
Support convolution in roofline.
0c8e601f
[Pallas] Add legacy PRNG key support to Pallas PRNG
dbd8d920
Explicit sharding docs
e8f43d1c
Merge pull request #27165 from jax-ml:sharding-in-types-doc
a11d8891
Expose `get_abstract_mesh` via the `jax.sharding` namespace
aa9480a4
Add `experimental/serialize_executable.py` to `BUILD`.
39e8ee93
Fix notebook formatting
412b2e3a
[Pallas] Increase #rows when casting to x2.
21f5f2d4
Merge pull request #27173 from jakevdp:fix-ipynb
95791fa9
[direct-linearize] shmap fixes
174dcc77
[pallas:mosaic_gpu] WG lowering now supports `while_p`
64230d1c
Merge pull request #27015 from mattjj:direct-linearize-fixes-4
b00a3a19
Merge remote-tracking branch 'origin/rocm-main' into ci-upstream-sync…
e275d5cf
Trivial change for CI
7a6940bb
charleshofer
approved these changes
on 2025-03-14
Login to write a write a comment.
Login via GitHub