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

26145 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
jax authors
aceae84fab [Pallas] Enable skipping of floating-point operations when interpreting Pallas TPU kernels on CPU. 
PiperOrigin-RevId: 735527650
 2025-03-10 15:14:00 -07:00
Jacob Burnim
802cb33bf8 [Pallas] Increase tolerance in PallasOutOfBoundsInterpretTest. 
PiperOrigin-RevId: 735519526
 2025-03-10 14:49:34 -07:00
jax authors
d55879723e Merge pull request #26840 from rajasekharporeddy:testbranch1 
PiperOrigin-RevId: 735513976
 2025-03-10 14:33:14 -07:00
jax authors
b8590816bf Merge pull request #26839 from Sai-Suraj-27:fix_jax.debug.print 
PiperOrigin-RevId: 735511953
 2025-03-10 14:26:45 -07:00
Sharad Vikram
81dde225b0 [Pallas/Fuser] Add select_n push rule 
PiperOrigin-RevId: 735510713
 2025-03-10 14:23:01 -07:00
jax authors
261e6e5fdc Merge pull request #27038 from jakevdp:vmap-sentinel 
PiperOrigin-RevId: 735510065
 2025-03-10 14:21:11 -07:00
jax authors
c942b0fef0 Merge pull request #26977 from jakevdp:fix-expn 
PiperOrigin-RevId: 735506133
 2025-03-10 14:09:32 -07:00
Sharad Vikram
87272fbe93 [Pallas/Fuser] Add debug option to fuser.fuse that prints out jaxpr 
PiperOrigin-RevId: 735505460
 2025-03-10 14:07:26 -07:00
jax authors
affe2e734e Rename dot_with_no_batch_dims_saveable to dots_with_no_batch_dims_saveable for internal consistency 
PiperOrigin-RevId: 735484326
 2025-03-10 13:04:49 -07:00
Praveen Narayanan
b6d4fe5387 Define lax.ragged_dot_general and express lax.ragged_dot in terms of it. 
PiperOrigin-RevId: 735471245
 2025-03-10 12:25:22 -07:00
jax authors
64beebbfb0 Merge pull request #27035 from vfdev-5:add-file-zip-to-tsan-ci-jobs 
PiperOrigin-RevId: 735467348
 2025-03-10 12:12:49 -07:00
jax authors
18f2f19c1a Merge pull request #26525 from wenscarl:e2m1fn 
PiperOrigin-RevId: 735457804
 2025-03-10 11:46:18 -07:00
Jacob Burnim
73d20cd62a [Pallas] Small fix to TPU interpret mode (input_output_aliases + scalar args). 
PiperOrigin-RevId: 735455671
 2025-03-10 11:40:10 -07:00
Jake VanderPlas
8ecadfdf9d Internal: make it easier to detect the vmap sentinel 2025-03-10 11:37:50 -07:00
jax authors
007fc7a6f1 Remove version limit for setuptools dependency. 
PiperOrigin-RevId: 735453796
 2025-03-10 11:36:17 -07:00
Nitin Srinivasan
d41e96835b Modify version test to consider "rc" versions as well 
I was testing the RC promotion workflow and found that the version test failed as it does not consider pre-releases. Therefore, this commit modifies the `VERSION_PATTERN` to also consider "rc" wheels.

Fixes https://github.com/jax-ml/jax/actions/runs/13705984545/job/38331236497

PiperOrigin-RevId: 735444828
 2025-03-10 11:10:18 -07:00
Michael Whittaker
5cb29949d4 Warn the user if transparent huge pages aren't enabled. 
PiperOrigin-RevId: 735431881
 2025-03-10 10:37:58 -07:00
jax authors
14b215fe76 Merge pull request #27032 from dfm:lax-dtype 
PiperOrigin-RevId: 735424674
 2025-03-10 10:18:58 -07:00
vfdev
1bab037ca0 Add file and zip to tsan.yaml 2025-03-10 16:51:05 +00:00
jax authors
ab0ce8a448 Merge pull request #26811 from dfm:direct-lin 
PiperOrigin-RevId: 735388827
 2025-03-10 08:39:49 -07:00
Dimitar (Mitko) Asenov
d2bf034c47 [Mosaic GPU] Test the wgmma_op lowering when a is in registers. 
I had to add support for wgmma layout in vector_load. Not sure if this is useful outside the test.

PiperOrigin-RevId: 735384104
 2025-03-10 08:25:43 -07:00
jax authors
5a7ef40416 Merge pull request #27026 from garymm:patch-3 
PiperOrigin-RevId: 735382490
 2025-03-10 08:20:30 -07:00
Dan Foreman-Mackey
21884d4a14 Move (most) jaxlib linalg custom call registration into JAX. 
My motivation here is to fix the plugin support for batch partitionable custom calls. Since plugin support for custom call partitioners is provided via register_plugin_callback in xla_bridge, instead of xla_client itself, it's much more straightforward to register the custom calls in JAX.

It would be possible to refactor things differently, but it actually seems like a reasonable choice to use the supported APIs from `jax.ffi` instead of `xla_client` so that we can take advantage of any new features we might add there in the future.

This is all still a little bit brittle and I'd eventually like to migrate to a version where the XLA FFI library provides a mechanism for exporting handlers, but this change is still compatible with any future changes like that.

PiperOrigin-RevId: 735381736
 2025-03-10 08:17:44 -07:00
Dan Foreman-Mackey
4eada56027 Avoid using array operations within lax.py operations. 2025-03-10 11:04:32 -04:00
Sergei Lebedev
91340ea0a7 [pallas:mosaic_gpu] Added support for math functions to the WG lowering 
PiperOrigin-RevId: 735333893
 2025-03-10 05:08:19 -07:00
jax authors
f906d2b6d1 Update XLA dependency to use revision 
efb27eb924.

PiperOrigin-RevId: 735322004
 2025-03-10 04:14:33 -07:00
Benjamin Chetioui
75d8702023 [Pallas/Mosaic GPU] Add lowerings/layout inference for all the necessary conversion ops when using Warpgroup semantics. 
Enable some of the pre-existing Pallas `ops_test`s for testing.

PiperOrigin-RevId: 735293084
 2025-03-10 02:14:39 -07:00
Gary Miguel
6a718b762f
Update stateful-computations.md 
tree_map -> tree.map
 2025-03-09 21:35:46 -07:00
jax authors
b9fb69d1fc Update XLA dependency to use revision 
e89a6b46bc.

PiperOrigin-RevId: 735080797
 2025-03-09 03:55:03 -07:00
jax authors
922935a916 Merge pull request #27006 from dfm:more-direct-lin-debug-info 
PiperOrigin-RevId: 734884478
 2025-03-08 06:46:44 -08:00
Dan Foreman-Mackey
36d515ed2c A few more fixes for debug_info tests with direct_linearize. 2025-03-08 07:47:24 -05:00
jax authors
04696b4d7b Update XLA dependency to use revision 
be68e80894.

PiperOrigin-RevId: 734851053
 2025-03-08 03:27:16 -08:00
Jevin Jiang
0f0636afab [Mosaic TPU][Pallas] Add pl.reciprocal 
PiperOrigin-RevId: 734749577
 2025-03-07 18:29:30 -08:00
jax authors
4988adccf1 Merge pull request #27010 from mattjj:direct-linearize-fixes-3 
PiperOrigin-RevId: 734747001
 2025-03-07 18:15:02 -08:00
Matthew Johnson
fe26c19b92 [direct-linearize] fix name_stack bugs 
Surprisingly, the bug was tracked down to #26111 aka cl/730939406, specifically
the new implementation of reset_name_stack in source_info_util.py.

To repro, use the before-this-commit implementation of reset_name_stack (left
commented-out in the file), and run

```
  JAX_USE_DIRECT_LINEARIZE=1 python tests/name_stack_test.py NameStackTransformationTest.test_nested_jit_stack
```
 2025-03-08 01:51:19 +00:00
jax authors
4660d7b6dd Merge pull request #27005 from mattjj:direct-linearize-fixes-2 
PiperOrigin-RevId: 734736244
 2025-03-07 17:17:45 -08:00
Matthew Johnson
251b93ebd7 fixups that we meant to include in #26427 
Co-authored-by: Dougal Maclaurin <dougalm@google.com>
 2025-03-08 00:03:26 +00:00
Jevin Jiang
041f575747 Support MHA in ragged paged attention for packed type 
PiperOrigin-RevId: 734695213
 2025-03-07 14:47:04 -08:00
jax authors
6095af050f Merge pull request #26427 from mattjj:direct-linearize-fixes 
PiperOrigin-RevId: 734687601
 2025-03-07 14:22:16 -08:00
jax authors
d849779689 Merge pull request #27001 from mattjj:yash-scan 
PiperOrigin-RevId: 734685031
 2025-03-07 14:14:30 -08:00
jax authors
1870176eb3 Merge pull request #26979 from mattjj:26936 
PiperOrigin-RevId: 734674945
 2025-03-07 13:43:55 -08:00
Matthew Johnson
f4f31f89ae [scan] when num_trips==0, don't generate weird size-zero reshapes 2025-03-07 21:35:40 +00:00
Matthew Johnson
7c2f842353 shard_map and other fixes to direct-linearize 
Co-authored-by: Dougal Maclaurin <dougalm@google.com>
 2025-03-07 21:02:40 +00:00
Matthew Johnson
0e30a3ace9 [mutable-arrays] read values should have the same explicit sharding as ref 
fixes #26936
 2025-03-07 20:53:29 +00:00
Hyeontaek Lim
178278863d [JAX] Fix api_benchmark broken by https://github.com/jax-ml/jax/pull/26569 
`pjit_check_aval_sharding` expects `names: Sequence[str]`.

PiperOrigin-RevId: 734614264
 2025-03-07 10:49:53 -08:00
jax authors
ccf7278292 Add the len(arg) to the error message for static_argnums 
Helps reduce the confusion on what is considered an argnum.
Ideally there should be static_argkwg

PiperOrigin-RevId: 734591856
 2025-03-07 09:49:49 -08:00
Yash Katariya
9f37b5197f [sharding_in_types] Fix a bug where empty_array in scan was created with the wrong spec when unroll > 1. 
PiperOrigin-RevId: 734591110
 2025-03-07 09:47:32 -08:00
Christos Perivolaropoulos
eeccc67c0b [mgpu] Debug print arrays. 
PiperOrigin-RevId: 734576543
 2025-03-07 08:58:25 -08:00
Adam Paszke
1bef8b61af [Mosaic GPU] Add a better explanation for the transposed layout 
Thanks to @bchetioui for the discussion!

PiperOrigin-RevId: 734564672
 2025-03-07 08:19:32 -08:00
Adam Paszke
402389290c [Mosaic TPU] Enable all conversions involving fp8 types on TPUv5+ 
PiperOrigin-RevId: 734558364
 2025-03-07 07:59:31 -08:00