mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 12:26:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
26,287 Commits 504 Branches 263 Tags
Commit Graph

16405 Commits

Author SHA1 Message Date
Jacob Burnim
47e8effdce Adds option to initialize buffers to NaNs or zeros in TPU interpret mode. 2025-03-18 12:24:45 -07:00
Yash Katariya
a5c0f200e7 set_mesh should return the prev_mesh instead of nothing. Users can choose to use the return value or ignore it. 
PiperOrigin-RevId: 738039559
 2025-03-18 09:43:25 -07:00
jax authors
7c5871f464 [Pallas TPU] Hoist prologue and epilogue outside of pipeline loop 
PiperOrigin-RevId: 738038138
 2025-03-18 09:40:43 -07:00
jax authors
30941480a1 Merge pull request #27198 from jakevdp:lax-docs 
PiperOrigin-RevId: 738038116
 2025-03-18 09:38:58 -07:00
jax authors
13541e9f12 Make blocked_fold_in consistent when the block sizes induce padding 
Add coverage for padded shapes to unit tests.

PiperOrigin-RevId: 738029476
 2025-03-18 09:12:11 -07:00
Jake VanderPlas
8b46e53a4f jax.lax: improve docs for several APIs 2025-03-18 08:55:38 -07:00
Benjamin Chetioui
1e36cbe597 [Mosaic GPU] Raise a NotImplementedError if swizzle=16. 
Unswizzled MMAs don't lower correctly, and are not currently intended to be
supported.

PiperOrigin-RevId: 737981373
 2025-03-18 06:29:13 -07:00
Adam Paszke
8da93249d2 [Mosaic GPU] Fuse slicing into s4 -> bf16 upcasts 
This allows us to significantly simplify the generated PTX/SASS,
which is currently cluttered with LLVM trying to align slices to
start at bit 0 and failing to CSE the right shifts.

PiperOrigin-RevId: 737967890
 2025-03-18 05:38:49 -07:00
Benjamin Chetioui
ba2f7c9ad9 [Mosaic GPU] Add transform inference rule for mgpu.slice_smem. 
PiperOrigin-RevId: 737957778
 2025-03-18 04:53:54 -07:00
Adam Paszke
d4bd2570ae [Mosaic GPU] Add a specialized layout for loading 4-bit inputs in WGMMA friendly layouts 
PiperOrigin-RevId: 737956598
 2025-03-18 04:47:51 -07:00
Adam Paszke
34cd5b0d74 [Mosaic GPU] Remove sub-byte conversion restriction 
XLA:GPU recently changed its endianness to little endian to better match LLVM
and the rest of the CUDA ecosystem, so we can lift the earlier restrictions.
PiperOrigin-RevId: 737934373
 2025-03-18 03:13:21 -07:00
Yash Katariya
549973dec6 Allow pspec to be passed to device_put if there is a mesh in the surrounding context 
PiperOrigin-RevId: 737812111
 2025-03-17 17:47:56 -07:00
Emily Fertig
8c35191725 Enable jax.device_put to a sharding with no local devices. 
PiperOrigin-RevId: 737797815
 2025-03-17 16:49:46 -07:00
Sergei Lebedev
051687dc4c [pallas] pallas_call_p is now parameterized by a mesh 
The mesh is necessary to add support for clusters to the Mosaic GPU backend.

PiperOrigin-RevId: 737792129
 2025-03-17 16:30:40 -07:00
jax authors
b4966130a3 Compute tile index using tile-based coordinates 
This reduces the chances of overflowing a 32-bit integer when computing tile indices.
Add unit test to reproduce the overflow with the previous implementation of `blocked_fold_in`.

PiperOrigin-RevId: 737778853
 2025-03-17 15:46:27 -07:00
Peter Hawkins
20658fabb3 Replace cached function get_replicated_hlo_sharding() with a constant. 
Small cleanup, no functional changes intended.

PiperOrigin-RevId: 737727727
 2025-03-17 13:17:33 -07:00
jax authors
ebcae0d30a Merge pull request #26980 from carlosgmartin:categorical_replace 
PiperOrigin-RevId: 737720590
 2025-03-17 12:58:01 -07:00
Peter Hawkins
be5d13af77 Remove code that preserved _original_py_fns on C++ classes. 
This no longer appears to be used.

PiperOrigin-RevId: 737715578
 2025-03-17 12:43:04 -07:00
Benjamin Chetioui
9a686e0bf3 [Mosaic GPU] Add initial transform inference rules for vector.{load,store}. 
PiperOrigin-RevId: 737703568
 2025-03-17 12:08:07 -07:00
carlosgmartin
3f59fa6888 Add replace option to random.categorical to enable sampling without replacement. 2025-03-17 13:41:46 -04:00
jax authors
de9ad6bad9 Merge pull request #27157 from mar-muel:improve-random-choice-performance 
PiperOrigin-RevId: 737665351
 2025-03-17 10:30:15 -07:00
Adam Paszke
3649da56fc [Mosaic GPU] Make the s4 -> bf16 upcast more flexible when it comes to vector length 
We can now perform the conversion in groups of 2, 4 or even 8 elements at a time.

PiperOrigin-RevId: 737626600
 2025-03-17 08:37:17 -07:00
Sergei Lebedev
0ff234049b Removed trivial docstrings from JAX tests 
These docstrings do not make the tests any more clear and typically just duplicate the test module name.

PiperOrigin-RevId: 737611977
 2025-03-17 07:49:37 -07:00
Sergei Lebedev
a7e5eaee56 [pallas:mosaic_gpu] jnp.reduce_sum now works for >1D arrays 
PiperOrigin-RevId: 737578598
 2025-03-17 05:32:07 -07:00
Adam Paszke
89b21de62a [Mosaic GPU] Add support for changing the layout before the upcast 
This lets us save on 2 ALU instructions (3x select becomes 1x prmt).

PiperOrigin-RevId: 737550598
 2025-03-17 03:26:48 -07:00
Adam Paszke
2bdd9c8797 [Mosaic GPU] Add support for fast WGMMA layout changes after 8- to 16-bit upcast 
PiperOrigin-RevId: 737542885
 2025-03-17 02:52:16 -07:00
jax authors
761b35c59e Merge pull request #27176 from jakevdp:lax-docs 
PiperOrigin-RevId: 737338493
 2025-03-16 05:39:55 -07:00
Joan Puigcerver
466ef6a132 Change the way that batching.spec_types is updated. 
There's no reason why not two custom vmappable types cannot share the same spec_type. However, spec_types was a set, which can cause bugs / exceptions.

Suppose that I register two vmappable data_types sharing the same spec_type, and then unregister one of the two. Then, the spec_type is no longer in the set to support the second data_type. Also, an exception will be raised if I try to unregister the two vmappable types (the second call to spec_types.remove).

When unregistering a data type, instead of removing its spec_type from the set, we regenerate the set from the remaining vmappable types.

PiperOrigin-RevId: 737280270
 2025-03-15 22:58:44 -07:00
Jake VanderPlas
de8b0564ce Better docs for jax.lax add/sub/mul/div 2025-03-15 11:49:51 -07:00
Ayaka
9b0ace4a11 Support error checking in explicit mode 
PiperOrigin-RevId: 737051146
 2025-03-14 18:58:26 -07:00
jax authors
7db59cdcca Merge pull request #27174 from mattjj:opt-barrier-ad-rules 
PiperOrigin-RevId: 737040381
 2025-03-14 17:59:07 -07:00
Peter Hawkins
14cb7453f0 Add a C++ implementation of a toplogical sort. 
This is an exact port of the current Python implementation to C++ for speed.

I am being careful not to change the topological order we return in any way in this change, although we may do so in a future change.

PiperOrigin-RevId: 737014989
 2025-03-14 16:04:25 -07:00
Matthew Johnson
dadc68b6c1 add experimental lax.optimization_barrier autodiff rules 2025-03-14 22:40:55 +00:00
jax authors
b00a3a1986 Merge pull request #27015 from mattjj:direct-linearize-fixes-4 
PiperOrigin-RevId: 737003323
 2025-03-14 15:24:11 -07:00
Sergei Lebedev
64230d1c93 [pallas:mosaic_gpu] WG lowering now supports while_p 
PiperOrigin-RevId: 736996154
 2025-03-14 14:59:29 -07:00
Matthew Johnson
174dcc771a [direct-linearize] shmap fixes 2025-03-14 21:38:50 +00:00
Daniel Suo
39e8ee93b0 Add experimental/serialize_executable.py to BUILD. 
PiperOrigin-RevId: 736975882
 2025-03-14 13:54:39 -07:00
Yash Katariya
aa9480a441 Expose get_abstract_mesh via the jax.sharding namespace 
PiperOrigin-RevId: 736972976
 2025-03-14 13:45:32 -07:00
Justin Fu
dbd8d92075 [Pallas] Add legacy PRNG key support to Pallas PRNG 
PiperOrigin-RevId: 736949584
 2025-03-14 12:30:04 -07:00
Zac Mustin
0c8e601f90 Support convolution in roofline. 
So far we support only `unfused_hmb_bytes` and don't account for `{feature, batch}_group_count`s due to complexity.

PiperOrigin-RevId: 736948528
 2025-03-14 12:26:20 -07:00
Yash Katariya
88d4bc3d45 Rename AxisTypes enum to AxisType 
PiperOrigin-RevId: 736935746
 2025-03-14 11:48:21 -07:00
Emily Fertig
bdb6d03322 Allow make_array_from_callback to construct nonaddressable arrays. 
PiperOrigin-RevId: 736922870
 2025-03-14 11:10:32 -07:00
Martin Muller
4a82fe94de Use lax.top_k instead of jnp.argsort in Gumbel top-k trick for weighted sampling without replacement in jax.random.choice 2025-03-14 19:02:24 +01:00
Ilya Tikhonovskiy
c9ac82c826 [XLA:GPU] Add missing BF16_BF16_F32_X9 matmul option in config.py 
Extend the list of possible default algorithms that dot could use.

PiperOrigin-RevId: 736879149
 2025-03-14 08:58:59 -07:00
Peter Hawkins
6fa98fc0a4 Use "x is y" rather than "id(x) == id(y)". 
The latter involves at least two object constructions.

PiperOrigin-RevId: 736878098
 2025-03-14 08:54:46 -07:00
Peter Hawkins
8ab33669e2 Add a variant of safe_map() that has no return value, named foreach(). 
This avoids a bunch of list bookkeeping in cases where we are iterating only for a side effect and do not care about the results.

I would have named this iter() to match OCaml's list.iter(), but unfortunately iter() is a Python builtin.

PiperOrigin-RevId: 736859418
 2025-03-14 07:42:48 -07:00
Peter Hawkins
074216e07a Precompute a weakref to a Trace≥ 
We use Trace weakrefs frequently, so we may as well construct one eagerly.

PiperOrigin-RevId: 736841778
 2025-03-14 06:26:17 -07:00
Benjamin Chetioui
5098d2ef49 [Mosaic GPU][NFC] Simplify implementation for in_{layout,transforms}_for_operand utils. 
PiperOrigin-RevId: 736809960
 2025-03-14 03:52:10 -07:00
Ilya Tikhonovskiy
43b78c539f [JAX] Add missing preset for X9 dot optimization on BF16/BF16 -> F32. 
Two PRs that support this feature have been submitted to stablehlo and openxla.
Now we could do the last step - enable it in JAX.

PiperOrigin-RevId: 736799241
 2025-03-14 02:57:55 -07:00
jax authors
cbece0b00b Add explicit support for float8_e4m3b11fnuz in pl.dot 
PiperOrigin-RevId: 736798315
 2025-03-14 02:51:55 -07:00