CI: 03/11/25 upstream sync #271
github-actions
merged
129 commits into
rocm-main
from
ci-upstream-sync-142_1
fix unit tests to not use fmha rewriter
3a800803
fix member access to packed CUDA struct
91cae595
fix Initializer protocol
5acfc88a
Fixed printing order of results in jax.debug.print documentation.
56285aec
Remove duplicate JAX version 0.4.37 heading in changelog
9c18e8dc
Improve error messages when input argument resolution fails in custom…
8b1b039e
Merge pull request #26765 from Qwlouse:patch-1
97db925a
Fix detection of epath
52ab8c4c
Update array-api-tests to latest commit
f0bbd26d
Merge pull request #26912 from dfm:resolve-args-error-message
4a73134b
Merge pull request #26916 from jakevdp:update-array-api-tests
7d0aab5a
[JAX] Disable msan and asan for the profiler test running on nvidia gpu
8af6f70f
jax.numpy ndim/shape/size: deprecate non-array input
8cec6e63
Internal config change
d112c85e
[Mosaic GPU] Add lowering for `log`, and a fast path using log2.
43b6be0e
Update setup.py to automatically pick up libtpu patch releases
bb80a568
Prepare for JAX release 0.5.2 (patch release over 0.5.1)
ce224293
Add functionality to allow promoting RC wheels during release
721d1a32
Merge pull request #26901 from NeilGirdhar:etils
b238bad7
Merge pull request #26641 from jakevdp:jnp-ndim
c145102e
Update all uses of `@tsl//third_party` to `@xla//third_party`
1a19d559
Update references to JAX's GitHub repo
cdeeacab
Make sure concat + vmap of sharded input and replicated input works p…
766315f7
Merge branch 'release/0.5.2' into main
a6c858f0
Update version number after 0.5.2 release
cebedb9f
[mgpu] Non-vector untiled stores for tiling layouts.
51719a1a
Update XLA dependency to use revision
06b760ee
Attempt 2 at landing custom_vjp.optimize_remat using custom_dce.
342cb7b9
Merge pull request #26627 from Cjkkkk:remove_fmha_rewriter
f3b2c841
Remove a TSAN suppression.
40e1a2a5
Merge pull request #26691 from h-vetinari:packed
a13b3ced
Removed unused import
6230ef1d
Merge pull request #26938 from superbobry:maint-2
9c19afd9
[Mosaic GPU][NFC] Refactor MMA SMEM descriptor creation
d1191387
Merge pull request #26934 from hawkinsp:tsan
4e1f969a
[Mosaic GPU] Add support for small tiles for (WG)MMA LHS
4493889c
support e2m1fn
c099e808
Reverts 342cb7b99a09180472823a33c7cdad8a8db77875
4a93c8b3
Update `.bazelrc` with Apple CC toolchain changes.
1ae7dd7f
[Mosaic GPU] Remove support for large tiles on Blackwell
8df00e26
add pmap axes hints
3e4dc0d4
Fix ambiguous cpu definition for JAX wheels.
3edc068f
vmap mismatch size error message: handle *args
69d66f66
Fix build rule for free-threaded python builds.
0913cd75
Merge pull request #26952 from garymm:vmap-arg
8571ad9f
[Pallas] Adds a simple dynamic race detector for TPU interpret mode.
016b351f
Set `USERPROFILE` for Windows builds to fix CI issue.
c16f37d8
Add a note about uneven sharding and with_sharding_constraint. Fixes …
ba5349f8
Just use `jit` as the string in error messages instead of `jit` and `…
a67ab9fa
[Pallas/Mosaic GPU] Enable `ops_test` for Mosaic GPU.
fe577b5d
Update XLA dependency to use revision
16bb9190
[pallas] Add support for `pl.dot` with `int8` inputs.
d6b97c20
[pallas:mosaic_gpu] Added WG lowering rule for `lax.bitcast_convert_t…
2a340193
Build JAX wheels instead of installing it from the source repository
623865fe
Minor bug fixes in error checking
8c89da7c
[Mosaic GPU] Fix `scf.ForOp` lowering to put lowered ops at the right…
5d64b3d2
Open source TPU-friendly ragged paged attention kernel.
4b49c035
Prevent tracer leaks in scipy.special.expn
b441b2b7
Updates the Colocated Python's serialization (and deserialization) im…
cd7f03f2
Merge pull request #26927 from skye:merge_release
4cab1183
Auto complete specs in a sharding if aval.ndim > len(sharding.spec) w…
e9486920
[Mosaic TPU] Support fp8 upcast to f32
ff4310f6
roofline: Support computing flops for unary ops.
8095d842
Fix lint
ccbe9f7c
Update XLA dependency to use revision
bf95bf49
Dedup non-ref constants closed in cond branch functions.
e6db7a9d
Merge pull request #26950 from lockwo:Owen/add-pmap-typehint
de78d2cc
[Mosaic GPU] Add support for tiling stores to refs using small tiling
85c6b6a1
Update ad.backward_pass to support non-linear functions of constants.
b7ecfdfd
Add a divisibility check so that we make sure that sharding evenly di…
f8b98993
[Mosaic GPU] Add a new layout to help with transposing WGMMA results
65462fe6
[pallas:mosaic_gpu] `copy_smem_to_gmem` now allows skipping `cp.async…
928caf83
[Mosaic TPU] Enable all conversions involving fp8 types on TPUv5+
40238929
[Mosaic GPU] Add a better explanation for the transposed layout
1bef8b61
[mgpu] Debug print arrays.
eeccc67c
[sharding_in_types] Fix a bug where `empty_array` in scan was created…
9f37b519
Add the len(arg) to the error message for static_argnums
ccf72782
[JAX] Fix api_benchmark broken by https://github.com/jax-ml/jax/pull/…
17827886
[mutable-arrays] read values should have the same explicit sharding a…
0e30a3ac
shard_map and other fixes to direct-linearize
7c2f8423
[scan] when num_trips==0, don't generate weird size-zero reshapes
f4f31f89
Merge pull request #26979 from mattjj:26936
1870176e
Merge pull request #27001 from mattjj:yash-scan
d8497796
Merge pull request #26427 from mattjj:direct-linearize-fixes
6095af05
Support MHA in ragged paged attention for packed type
041f5757
fixups that we meant to include in #26427
251b93eb
Merge pull request #27005 from mattjj:direct-linearize-fixes-2
4660d7b6
[direct-linearize] fix name_stack bugs
fe26c19b
Merge pull request #27010 from mattjj:direct-linearize-fixes-3
4988adcc
[Mosaic TPU][Pallas] Add pl.reciprocal
0f0636af
Update XLA dependency to use revision
04696b4d
A few more fixes for debug_info tests with direct_linearize.
36d515ed
Merge pull request #27006 from dfm:more-direct-lin-debug-info
922935a9
Update XLA dependency to use revision
b9fb69d1
Update stateful-computations.md
6a718b76
[Pallas/Mosaic GPU] Add lowerings/layout inference for all the necess…
75d87020
Update XLA dependency to use revision
f906d2b6
[pallas:mosaic_gpu] Added support for math functions to the WG lowering
91340ea0
Avoid using array operations within lax.py operations.
4eada560
Move (most) jaxlib linalg custom call registration into JAX.
21884d4a
Merge pull request #27026 from garymm:patch-3
5a7ef404
[Mosaic GPU] Test the wgmma_op lowering when a is in registers.
d2bf034c
Merge pull request #26811 from dfm:direct-lin
ab0ce8a4
Add file and zip to tsan.yaml
1bab037c
Merge pull request #27032 from dfm:lax-dtype
14b215fe
Warn the user if transparent huge pages aren't enabled.
5cb29949
Modify version test to consider "rc" versions as well
d41e9683
Remove version limit for `setuptools` dependency.
007fc7a6
Internal: make it easier to detect the vmap sentinel
8ecadfdf
[Pallas] Small fix to TPU interpret mode (input_output_aliases + scal…
73d20cd6
Merge pull request #26525 from wenscarl:e2m1fn
18f2f19c
Merge pull request #27035 from vfdev-5:add-file-zip-to-tsan-ci-jobs
64beebbf
Define lax.ragged_dot_general and express lax.ragged_dot in terms of it.
b6d4fe53
Rename `dot_with_no_batch_dims_saveable` to `dots_with_no_batch_dims_…
affe2e73
[Pallas/Fuser] Add debug option to fuser.fuse that prints out jaxpr
87272fbe
Merge pull request #26977 from jakevdp:fix-expn
c942b0fe
Merge pull request #27038 from jakevdp:vmap-sentinel
261e6e5f
[Pallas/Fuser] Add select_n push rule
81dde225
Merge pull request #26839 from Sai-Suraj-27:fix_jax.debug.print
b8590816
Merge pull request #26840 from rajasekharporeddy:testbranch1
d5587972
[Pallas] Increase tolerance in PallasOutOfBoundsInterpretTest.
802cb33b
[Pallas] Enable skipping of floating-point operations when interpreti…
aceae84f
Better error message when `raise_if_error()` is called within a trace…
988a1208
[Pallas TPU] Remove `next_slot` SMEM tensor from pipeline emitter
02505fa7
Under pjit the `with mesh:` context will use `use_mesh(mesh): jit` in…
76dec382
PR #22800: Change the default value of print_operand_shape_ to false …
cb2eb157
Update XLA dependency to use revision
b6da46ec
[Pallas/Mosaic GPU] Explicitly disable `ops_test` on Mosaic GPU pre-H…
7fd32ecc
Merge branch 'rocm-main' into ci-upstream-sync-142_1
fb89a4b4
charleshofer
approved these changes
on 2025-03-11
Login to write a write a comment.
Login via GitHub