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

26292 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
Yash Katariya
76d9890bb7 Run the stream annotation tests on 2 devices so that it can be tested in TAP 
PiperOrigin-RevId: 738113725
 2025-03-18 13:01:48 -07:00
Matthew Johnson
942ff38e36 fix to ragged_all_to_all transpose 
PiperOrigin-RevId: 738110447
 2025-03-18 12:51:21 -07:00
Benjamin Chetioui
875099b25d [Mosaic GPU] Enable the new transform inference pass in the warpgroup lowering. 
A couple of dummy transform inference rules needed to be added in order to
contend with parts of the lowering that do not use the dialect yet, along with
a transform inference rule for `memref.view`.

PiperOrigin-RevId: 738089782
 2025-03-18 11:51:43 -07:00
Peter Hawkins
547d602760 Remove //jaxlib:cpu_kernels and //jaxlib:gpu_kernels forwarding Bazel targets. 
These were temporary forwarding targets that are no longer needed; use //jaxlib/cpu:cpu_kernels and //jaxlib/cuda:cuda_gpu_kernels instead.

PiperOrigin-RevId: 738085234
 2025-03-18 11:39:00 -07:00
jax authors
ee0073e605 Merge pull request #27094 from vfdev-5:fix-tsan-numpy-install-patch 
PiperOrigin-RevId: 738080051
 2025-03-18 11:24:15 -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
vfdev-5
9145d617b8 Added exit 1 if git patch is failed + other checks 2025-03-18 16:47:34 +01: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
jax authors
7a459f0ed1 Update XLA dependency to use revision 
3bb7654721.

PiperOrigin-RevId: 737959582
 2025-03-18 05:01:30 -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
Chris Jones
38d52a19ef [mosaic_gpu] Force flush all cupti activity, then unsubscribe. 
With default flushing, it is possible for events to be missed. We should only unsubscribe after we are finished with cupti.

PiperOrigin-RevId: 737939327
 2025-03-18 03:35:03 -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
Changhui Lin
f174b00f23 Replace the uses of PjRtClient::Compile() with PjRtClient::CompileAndLoad(). 
This is to prepare for updating `PjRtClient::Compile()` to return an unloaded executable [1/N]

PiperOrigin-RevId: 737805623
 2025-03-17 17:18:31 -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
jax authors
b74b16f9b9 Merge pull request #27164 from MichaelHudgins:a4-testing 
PiperOrigin-RevId: 737733904
 2025-03-17 13:36:04 -07:00
Michael Hudgins
ecf7fde714 Add B200 testing to continuous workflow 2025-03-17 20:19:20 +00:00
jax authors
4f70471310 Fix error in pallas tutorial 
PiperOrigin-RevId: 737727935
 2025-03-17 13:19:12 -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
Nitin Srinivasan
031614c22b Pin numpy~=2.1.0 in workflow file instead of test-requirements.txt 
PiperOrigin-RevId: 737632771
 2025-03-17 08:59:06 -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
jax authors
55812c5d02 Update XLA dependency to use revision 
fcf97e619e.

PiperOrigin-RevId: 737581187
 2025-03-17 05:44:17 -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
jax authors
e8b683aee0 Update XLA dependency to use revision 
936a727db7.

PiperOrigin-RevId: 737338103
 2025-03-16 05:37:58 -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
jax authors
f360e19194 Update XLA dependency to use revision 
f52d5e03ce.

PiperOrigin-RevId: 737143950
 2025-03-15 05:10:35 -07:00
Ayaka
9b0ace4a11 Support error checking in explicit mode 
PiperOrigin-RevId: 737051146
 2025-03-14 18:58:26 -07:00
jax authors
d07d642d6f Merge pull request #27177 from jax-ml:mixing_modes 
PiperOrigin-RevId: 737047069
 2025-03-14 18:34:27 -07:00
Yash Katariya
3c0027af3b mixing modes 2025-03-14 18:23:27 -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