[Test do not merge] Testing build wheels #722
alekstheod
wants to merge
7513 commits into
main
from
amd-main
![dependabot[bot]](https://avatars.githubusercontent.com/in/29110?v=4&s=60)
Bump actions/checkout from 6.0.0 to 6.0.1
523b667b
Propagate input sharding to output in full_like if it's fully replica…
3b6fd04e
[hijax] add is_vjp bit to LinearizeTrace
79f3519e
Merge pull request #35285 from mattjj:is-vjp-bit
4181fe49
Update XLA dependency to use revision http://github.com/openxla/xla/c…
4f22f9d2
Canonicalize indexer.slice_sharding. Fixes https://github.com/jax-ml/…
7fbbe019
[hijax] improved dce interface
b1981cbb
Merge pull request #35301 from mattjj:better-hijax-dce-interface
7f80f430
[hijax] fix followup on #35301
acad3005
Merge pull request #35302 from mattjj:hijax-dce-fix
0cff9fe3
[hijax][remat3] start implementing remat on hijax
5f43a720
Merge pull request #35304 from mattjj:remat-transform
53d6b350
[hijax][remat3] add basic remat_jaxpr
b17e41dd
Merge pull request #35306 from mattjj:remat-jaxpr
dea31b6c
Update XLA dependency to use revision http://github.com/openxla/xla/c…
8dac162f
[pyrefly] Fixed or suppressed type errors in Pallas backends
9ce7295f
Merge pull request #35309 from superbobry:pyrefly
4d3395f0
Update XLA dependency to use revision http://github.com/openxla/xla/c…
f47f0c20
[pallas:mgpu] Fix handling of `AbstractRef` in `_commute_transform`.
cd019717
Merge pull request #33802 from jax-ml:dependabot/github_actions/actio…
a58b181f
Skip scaled matmul tests on CPU
11558844
[test] fix deprecated mpmath import
b5c92015
More Pyrefly fixes in Pallas
0beb8372
Merge pull request #35323 from jakevdp:mpmath-fix
adc7836d
Merge pull request #35317 from superbobry:pyrefly
ad8c8c4f
[typing] more pyrefly fixes
b1c5e118
Another batch of Pyrefly fixes
e673505c
[Mosaic TPU] Give more time for the changes to propagate to libtpu
db34816c
Fix dangling std::string_view in keys of GetModuleImage() cache
01c30830
[jaxlib] Bundle _{tpu,triton,mosaic_gpu}_ext.pyi with jaxlib
22875809
Pin mpmath to v1.3.0.
83471837
Merge pull request #35332 from superbobry:pyrefly
b486aacf
Fixed `self.accum_ref` assertion in mosaic/pipeline.py
4dc41c30
Roll back adding __eq__ and __hash__ to PyDevice.
920f6665
[pmap] Replace `functools.lru_cache` with `util.cache` to fix stale `…
fc5ab2c1
[lax_numpy] Adjust tolerance for np.float32 in trapezoid tests.
74b70ef1
Bump actions/setup-python from 6.1.0 to 6.2.0
ef24700b
Bump actions/cache from 4.3.0 to 5.0.3
9c33cd1c
Merge pull request #35344 from jax-ml:dependabot/github_actions/actio…
45e29f6f
[typing] suppress mypy redundant-cast
5d01b1fe
Merge pull request #35347 from jax-ml:dependabot/github_actions/actio…
8ac9b74c
Re-enable some disabled tests in pjit_test.py
cc7365fa
Merge pull request #35343 from jakevdp:fix-mypy-lint
677d7b0d
Merge pull request #35331 from jakevdp:more-pyrefly-fixes
e12cf3d5
[mlir] Fixed the return type of `*{Attr,Type}.get` staticmethods
6dc44197
Update XLA dependency to use revision http://github.com/openxla/xla/c…
d6d975c3
[Mosaic GPU] Add support for swizzle=16 in async SMEM->TMEM copies
b55efd81
Reverts 42c458cc318706bbbefbb45a752da6615a7123d3
776d3c38
More Pyrefly fixes
98f0ca20
[Pallas:MGPU] Allow SMEM->TMEM copies with unswizzled (but tiled) refs
1807a491
[Mosaic GPU] Relax N multiplicity restrictions for tcgen05 block-scal…
192601ce
[Mosaic GPU] Add support for fp4 tcgen05 MMAs that are both block-sca…
d037727f
Merge pull request #35133 from mwhittaker:fault_tolerance_blog_tweaks
c139db7a
[Mosaic TPU] Expose the latest version from the TPU MosaicSerdePass
ab5ef253
Add a light, blocking Mosaic B200 presubmit.
f0f74df5
[Pallas:MGPU] Add support for writes to WGMMA accumulator refs
fc9061d2
[Mosaic:GPU] Temporary fix of ASAN test.
e16cf237
PR #35345: Bump actions/upload-artifact from 5.0.0 to 6.0.0
afa13c42
Merge pull request #35367 from superbobry:pyrefly
012d6659
[mosaic] Verify that `tpu.memref_slice` always produces tile-aligned …
f59d8213
[hijax] mlir.lower_fun should have lo-jaxprs
b69d1c07
[NFC] Remove some unnecessary skips to have better test coverage.
e76e398b
Add Pallas SparseCore guide
6b89f7cb
Merge pull request #35379 from mattjj:mlir-lower-fun-to-lojax
2cbf65c1
Reverts f59d8213d42d124048de9ad93792bc17152134c7
7b56a865
[Mosaic TPU] Add more complete support for 16-bit `cmpf` on TPU v5. P…
518b0d11
Update XLA dependency to use revision http://github.com/openxla/xla/c…
0f31a1b0
Add SC support for pltpu.trace_value
2e3c8407
Merge pull request #35334 from ROCm:ci_arech_fix_dangling_string_view
f6af222b
[Pallas:MGPU] Add a test for tcgen05 MMA that's both sparse and block…
90fdba0f
[windows] Match triton importing in pallas_shape_poly_test to how it'…
eda5ec98
[jaxlib] Every MLIR dialect now comes with a .pyi
9169891b
[lax_scipy] Adjust tolerance for np.float32 in trapezoid tests.
ab5f1971
[lax_numpy] Skip ReducerWhere tests temporarily.
dc112482
[Mosaic TPU] Allow non-DMA semaphores in EnqueueDMA
ace7a790
Merge pull request #35357 from superbobry:more-piys
1ae8be13
[lax_numpy] Skip AdvancedIndexing tests temporarily.
122abbba
Enable some tests with TC-tiling with newer libtpu.
c6d2109c
Fix a bug in how input/output memory spaces are detected and propagat…
9c9dd77e
[typing] add type annotation for shard_map
709086fa
[Pallas:MGPU] Fix a race in dynamic_scheduling_loop
84a33c08
[ci] Update pyrefly pre-commit to v0.54.0
a211a318
[windows] Skip ops_test.test_binary_scalar if no TPU.
5c0ec3e8
[lax_numpy] Skip setxor1 tests temporarily.
ac537c76
Merge pull request #35404 from jakevdp:pyrefly-054
f5c39fed
Merge pull request #35015 from IvyZX:scguide
100b05c8
[pyrefly] remove ignore statements for issues fixed in v0.54
19e71a58
[typing] add type annotation for `jax.smap`.
b8171e35
Merge pull request #35405 from jakevdp:pyrefly-ignore-statements
cda29303
Merge pull request #34601 from ROCm:phambinh/ci_rocm_nn_tests
1f3139e0
Merge pull request #35178 from ezhulenev:jax_sort_devices_by_process_…
015f1cbd
[Pallas][Mosaic TPU] Add flag to disable semaphore checks per kernel.
77805318
[typing] more pyrefly fixes in jax/_src/lax/
0461a3bb
Merge pull request #35325 from jakevdp:pyrefly-fixes
742ad02a
Removed more dead code
480847f8
[pyrefly] fix remaining jax/_src/lax errors
1e102aeb
[windows] Skip triton export tests in pallas_shape_poly_test if trito…
7f250e0a
Explicitly check for disallowing TransformedRefs in higher-order JAX …
6a635c9a
[typing] pyrefly fixes in jax/_src/scipy/stats
b0d91345
[typing] pyrefly fixes in jax/_src/scipy/optimize
3b399a8f
[mosaic] Removed previously deprecated tpu.wait_dma
6869e28a
Merge pull request #35414 from levskaya:pyrefly3
dc3f50ca
Reverts 6869e28a4bde0f6d3a06a7c39858fe40cc44ff60
efa6bd9f
[Mosaic:GPU] Simplify collective metadata initialization.
6f06ab80
Update XLA dependency to use revision http://github.com/openxla/xla/c…
d0d6f426
Upgrade Abseil to LTS 20260107.1
ee4cea11
Merge pull request #35412 from levskaya:pyrefly
4826da47
[Mosaic GPU] Add support for warp-level f8 MMAs
7192cf7d
[Pallas:TPU] Add tests for DMAs with regular semaphores
9e11f1c0
[Pallas:TPU] Add a test checking that dynamic shape exports require t…
f2017f78
Set rpath based on the type (link_only)
af59a645
Add int1/uint1 dtypes.
c141c0fc
[mosaic] Another go at removing tpu.wait_dma
0b7d6672
[Mosaic GPU] Relax restrictions on N for sparse tcgen05 MMA
3f667fd4
Raise a better error if jit(shard_map) is bound with an AbstractMesh …
7af1783e
Fix hijax `VmapOf.batch_dim_rule` with `None` map dims.
6df7333b
[hijax] cleaner scan inc_rank/dec_rank api
4b87f031
[JAX] Bound tracebacks up to the enclosing `jit`.
6f2f0f36
Add a pre-commit hook that checks for Apache licence headers.
66b61418
Suppress still-broken test.
8f96afaa
Roll back skipped tests that were failing on LLVM bit/byte-packing is…
873bebe1
Update XLA dependency to use revision http://github.com/openxla/xla/c…
c3b53a26
Add jaxlib version guard for updated make_c_api_client.
17cfbb81
[jaxlib] Ported normalize_stubs.sh to Python
6428a91d
Reverts 7af1783eb619f7e0236aa8c25fe1fc72881bb6b3
12344fc0
Tie LinearizeTrace.tag to tangent_trace.tag
b6122922
Merge pull request #35439 from mattjj:hijax-cleaner-scan-axis-spec
3b482638
Merge pull request #35413 from jakevdp:pyrefly-lax-fixes
ed628980
[hijax] lower jaxpr during discharge_state
fa369541
[Pallas/TPU] Add tpu_info.get_tpu_info_for_chip
d237804a
Skip pallas lowering determinism testOrderAgnostic test if not on TPU.
f82e2810
[typing] miscellaneous pyrefly typing fixes
85f6ef56
[typing] fixed remaining scipy pyrefly typing issues
47291b9f
Generalize how op verifiers checking issuing core
008467a7
Fix undefined name in ad_checkpoint
e81437f2
Merge pull request #35219 from jakevdp:fix-tracer-tag
6d14b201
Merge pull request #35452 from jakevdp:fix-ad-checkpoint
c08a1707
Merge pull request #35448 from jakevdp:discharge-hijax
153937a3
[linux-arm64] Reduce parallelism to 32 for nogil pytest cpu to avoid …
d873c960
Merge pull request #35421 from levskaya:pyrefly4
019a87b1
[pyrefly] fix errors in jax/interpreters/ad
d57e7c8b
[pmap] Accelerate deprecation of PmapSharding.
11af98b4
Merge pull request #35422 from levskaya:pyrefly5
e008d9d8
Add an experimental `top_level_all_gather` API (private for now).
a6df1f29
Add a test for replicated -> sharded+unreduced reshard
d506abdb
Update XLA dependency to use revision http://github.com/openxla/xla/c…
08a51cea
Merge pull request #35102 from ROCm:set_release_rpaths_to_rocm_so_tar…
717f99a4
[Pallas:MGPU] Allow inline_mgpu inside a warp context
df6415b4
[Pallas][TPU kernel interpreter] Consistently use `pltu.HBM` to refer…
caf1324f
Update XLA dependency to use revision http://github.com/openxla/xla/c…
0d37f593
Skip pallas lowering determinism testOrderAgnostic test if not on TPU.
a573e865
[linux-arm64] Reduce parallelism to 32 for nogil pytest cpu to avoid …
969da1b7
Reverts ee4cea11e08b92f23f464be2e7795f5eb9fe9557
143c2366
Remove references to tsl/platform/logging.h header.
89469861
[pallas:sc] Moved `jaxpr_call` into core Pallas
6400fa1a
Fix typo in function parameter documentation
fd7c2aca
[mosaic] Added pattern canonicalizing tpu.memref_squeeze(memref.cast(…
3ad94844
[Mosaic GPU] fix annotation typo: s/pytype/pylint/
10a59a5e
[Mosaic TPU] Make sure to take the export flag into account when not …
b9926c39
[mosaic_gpu] Fixed mosaic-gpu-resolve-trivial-locations
e3adc6ad
[Mosaic:GPU] Re-use the barrier buffers between executions to prevent…
1609c18f
[Mosaic:GPU] Move collective kernel loading to the prepare stage.
f3bf01ad
[Mosaic:GPU] Add test-case for several mosaic ops.
a25a24df
Reimplement weakref_lru_cache on top of a custom hash map type.
b2ccb32c
[Mosaic TPU] Make sure to allow unregistered dialects before deserial…
0a513d30
Integrate LLVM at llvm/llvm-project@1053047a4be7
a30d6bc2
Skip tests that are failing at libtpu=0.0.35, but passing at libtpu-0…
8350f362
Bump jaxlib_extension_version to 412 due to LLVM integrate.
438e6946
Skip tests that are failing at libtpu=0.0.35, but passing at libtpu-0…
0eb23146
[Mosaic] Remove patterns to fold memref.cast into tpu.enqueue_dma and…
1899c328
Rename `pypi_latest` libtpu-version-type to `pinned` and use pinned l…
abd765d1
Add workflow for build wheels and running bazel tests on ROCm
bec3b790
Raise an error on `f32[4]{R:x} * f32[]`. For example `x: f32[4]{R:x} …
85ec15da
Rename `pypi_latest` libtpu-version-type to `pinned` and use pinned l…
70a8149d
Reverts b2ccb32cabd9da9c99a4a9b7f163d4059ff448d5
41ae1eb6
[Mosaic] Simplify GetCoreTypeOfParentOp
f154eb02
Fix release workflow by using the folder with pre-downloaded jax whee…
cfe83cb2
Fix release workflow by using the folder with pre-downloaded jax whee…
162b2ed0
[Pallas][Mosaic TPU] More complete 1D block support in Pallas/Mosaic.
5c96b9ee
[pallas] rewrite pull_block_spec rule to work in terms of block_index…
0b565503
Update XLA dependency to use revision http://github.com/openxla/xla/c…
b66a3c36
Skip layout_test.test_layout_donation_mismatching_in_and_out_fails fo…
e5ab8054
Skip layout_test.test_layout_donation_mismatching_in_and_out_fails fo…
7aa10a02
Skip sparse random and dot general sampled ad tests.
500bbc20
Skip sparse random and dot general sampled ad tests.
bb3362eb
Prepare for JAX release 0.9.1
58cb6e55
Update XLA dependency to use revision http://github.com/openxla/xla/c…
eba07b3e
Update XLA dependency to use revision http://github.com/openxla/xla/c…
b0d504a3
Add C addition to hopper matmul kernel.
d4204d58
Merge remote-tracking branch 'origin/release/0.9.1' into postrelease/…
29ce624a
Postrelease JAX v0.9.1.
5b9d90af
[jaxlib] :normalize_stubs now also fixes improper ``get`` classmethods
c5d55755
[mgpu] Minor cleanup in `try_cluster_cancel`.
95d95e3d
[Mosaic GPU][NFC] Remove duplicate layout inference rule for `vector.…
0cd1592e
Merge pull request #35502 from superbobry:stubgen
b4e692ea
[Mosaic GPU][NFC] Use more precise type for `op` in `_optimization_ba…
e20119d1
[pallas:sc] Switched to using `pltpu.emit_pipeline` in the lowering
41430d32
[Mosaic] Add assembly format to tpu.delay_op
31fc0bd5
[Mosaic GPU] Fix for the latest LLVM version
44736b92
Add serialization support for Mosaic GPU kernels.
a3ffef5e
Regenerate type stubs for nanobind v2.12.0.
0ded2fed
[Mosaic GPU][NFC] Remove dead `constraints.BroadcastInDim`.
ed6d58f2
Merge pull request #35501 from jax-ml:postrelease/0.9.1
edf82d10
Reverts 41ae1eb6ea50739feb91589705aa078b1495d40f
68b2565e
Merge pull request #34702 from ROCm:add-run-bazel-test-rocm-rbe
ba54b343
Merge pull request #35417 from joshuapjacob:patch-1
a192b434
[Mosaic GPU] Add support for untiled CP_ASYNC
a1b958ba
Merge pull request #35453 from jakevdp:pyrefly-ad
4b423a65
Remove nvidia_wheel_versions
2fe53184
Make jaxlib targets visible
c7ec794c
hipblas typedef fix
789df45d
No GPU fail
d0e32075
Wrap HIP inline functions in anonymous namespaces in vendor.h
caf56b60
SWDEV-512768 - Replace hipGetLastError with hipExtGetLastError
d8deb09f
Add shared utility function get_rocm_version to test_util.py
6eee40dd
Fix hipSparse CSR algorithm mappings for ROCm 7
3128e0e7
Fix v_pages quantization and adjust test params for ROCm compatibilit…
cfddfa60
Address LLVM assertion failure due to a multithreaded use. Update .gi…
cb904b38
Add skip of test_is_finite() on Cuda (#565)
b78fe9a6
Add rocm test requirements file (#570)
514fec32
Let the unit tests use build.py for setting up Bazel commands for uni…
b64a5e40
adding abort logic to rocm/jax (#590)
640e4a21
Skip is_finite tests on ROCm (not in Triton lowering for jax 0.8.0) (…
6eb4fdf5
Fix shared memory limit check for ROCm in test_dot (#596)
14b0125c
Fix Numpy signatures test (#598)
92211a5d
fix merge arts
b202e1f4
Enable RngShardingTests (#644)
4c60b2b0
Enable test_variadic_reduce_window on ROCm (#647)
22ac32aa
Skip sparse tests on ROCm due to hipSPARSE issue (#652)
8ecdd8b7
Update sparse test skip messages in v0.8.2 (#653)
cc543662
Skip sparse tests on ROCm due to hipSPARSE issue (#652)
1b6c9e72
Update sparse test skip messages in v0.8.2 (#653)
f0e90395
Skip sparse tests on ROCm due to hipSPARSE issue (#652)
0cc3b988
Update sparse test skip messages in v0.8.2 (#653)
47321964
Enable testMultivariateNormalSingularCovariance on ROCm (#666)
e5714928
Skip test_tridiagonal_solve on ROCm due to hipSPARSE numerical errors…
84ad452e
Update Skip Reason Outputs (#663)
3c295779
Skip sparse tests on ROCm due to hipSPARSE issue (#652)
7648e11b
Update sparse test skip messages in v0.8.2 (#653)
e0d2d1e9
Skip testCudaArrayInterfaceOnNonCudaFails on ROCm platform (#677)
a186d625
Skip sparse tests on ROCm due to hipSPARSE issue (#652)
9983dd00
Update sparse test skip messages in v0.8.2 (#653)
c45c49a9
Skip sparse tests on ROCm due to hipSPARSE issue (#652)
14d1b171
Update sparse test skip messages in v0.8.2 (#653)
ec394330
Add ROCm encoding for test_struct_encoding_determinism (#683)
d18c7571
Remove 'mean' from unsupported params for jnp.var (#689)
717cd20f
Implement approx_tanh for ROCm using OCML tanh function (#691)
02ae6c5b
Skipping testEighTinyNorm due to hipSolver issues (#697)
f2f9dd18
Abort detection CI workflow (#688)
6957e42d
Abort-Detection: Fix halt-for-connection input (#712)
7d684aa5
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Login to write a write a comment.
Login via GitHub