Merge amd-main into dpx_pytest #885
mkunredd
wants to merge
1179 commits into
dpx_pytest
from
amd-main
Increase CPU shard counts for `api_test`, `lax_numpy_hijax_test`, and…
0d8897f0
[pallas] DCE a pallas_call whose outputs are unused and whose Ref inp…
24040da5
Avoid unflattening None during side-effect pytree traversals.
1167501f
Merge pull request #41114 from mattjj:push-oszlxkxnouyw
61ca870a
Canonicalize constants in remat and fake_linear_op.
342ed50c
Merge pull request #40253 from Tixxx:tixxx/dynamic_root_pbroadcast
0c1de216
Allow IOEffect and OrderedIOEffect in custom_derivatives_allowed_effe…
4c71acb6
Fix `jax.scipy.special.erfcx` being 0 for certain ranges of `x`, and …
c12b7b40
arr.at[..].get(): support strategy argument
e38ed12b
Remove xla_tpu* flags for all_reduce.
59652f67
Merge pull request #41106 from stroxler:improve-random-f-validation
50479b03
Merge pull request #41112 from jakevdp:indexing-strategy
8398d4c2
internal change
099f4231
Bump jaxlib_extension_version to 501 for export_to_xprof and ProfileD…
09432f86
Add nary layout rule which errors out if all the operands don't have …
38b1ee0f
Support batched inputs in jax.scipy.linalg.lu
66ca0dce
[pallas:tpu] Migrate `dma_wait` lowering to `tpu.wait_dma`
58a76019
[pallas:sc] Allow specifying sem=None and wait_write() for streams on…
23e25acb
[internal] refactor to remove private array._chunk_iter
2ff6fdcf
[mgpu] Cleanup predicates in primitives.py.
3f78a820
[pallas:sc] Ignore the push stream warning in the trace value test
de32488f
[pallas:sc] Add src_sem and wait_read() for local async_copy on Spars…
41911d92
Update XLA dependency to use revision http://github.com/openxla/xla/c…
ff2913d5
[Mosaic GPU] Optimize iteration order in `ConstraintSystem.__and__` a…
18b7f840
[Mosaic GPU] Avoid reconstructing unchanged constraints in `reduce_co…
be2c8888
[TPU] Add regression test for https://github.com/jax-ml/jax/issues/38005
ab1a6295
[pallas:mosaic] Avoid recomputing the grid indices twice in `emit_pip…
49c942d7
Merge pull request #41129 from jakevdp:refactor-iter
e44dfa91
Merge pull request #40931 from LiRunGuo:fix-scipy-lu-batched
ecde8309
Update unary_ops_accuracy_test to account for improved expm1 accuracy…
3243e2b4
[pallas:mosaic] Avoid passing the deprecated `device_id_type=` in :tp…
847a0c33
[pallas:mosaic] Removed usages of the deprecated `device_id_type=` fr…
fed08424
jax.numpy: forward `initial` from Array.sum() to jnp.sum()
1038e1ee
Replace float0 cotangents with Zero in CustomVJPTraced.vjp_bwd_retval…
b4dae99a
Merge pull request #40877 from ROCm:zahiqbal/gfx1250-targets
7e161496
Merge pull request #41137 from inchang-ing:fix-41120-array-sum-initial
149c8e9c
[Mosaic TPU] Fix some broadcast -> reshape failures.
3cec6dde
Use pinned_host D2H transfers in array_serialization only on GPU.
f0d9ce76
[test] expand lax_numpy_reducers_test.py scaffold to test method-base…
3b32c93e
[Mosaic GPU] Bump the forward-compatible IR version to 7.
97b33c4d
Increase shard counts for `api_test` (GPU) and `svd_test` (CPU) to pr…
004eda6f
[custom_jvp] keep rules in partial eval, and don't force remat in cus…
620a3b41
Merge pull request #41152 from mattjj:push-qlrkmnspttmt
35e4da00
[Pallas SC] Enable vector layout passes across passing SparseCore uni…
f608ab39
Split elementary_test.py and special_test.py into grouped numerics te…
4c4ca825
jnp.atleast_nd: return tuples rather than lists
3a9db9c1
[XLA] Fix off-by-one boundary check for GT in WhileLoopSimplifier::Tr…
44e80cc9
Update XLA dependency to use revision http://github.com/openxla/xla/c…
6ef8456c
[Mosaic GPU] Bump the forward-compatible IR version to 8.
16a73c19
[jaxlib] Pre-register `PyDeviceList` iterator types during module ini…
06021174
Fix keyword argument forwarding in fwd_and_bwd
802ce557
Support generalized Hermitian eigenproblems in scipy.linalg.eigh
672c1247
Fix pareto support test when scale is below the resolution of loc
35c393e2
Fix unitarity check in LaxBackedScipyTests.testPolar for rank-deficie…
3bed0f84
[ROCm] Drop gfx1250-strict from the default ROCm targets
3104fab5
[Mosaic GPU] Use constraint provenance to improve layout inference er…
81a4cc91
[JAX] Cache lowered functions in `_lower_jaxpr_to_fun_cached` when `c…
03fd8a01
[oneapi] Set jax_oneapi_visible_devices in jax.distributed.initialize
9e83f074
[oneapi] Fix oneAPI CI GPU pinning and serialize multi-accelerator ba…
67788c4b
Merge pull request #41156 from jabrailkhalil:fix/fwd-and-bwd-keyword-…
9379227e
[JAX] Update SparseCore custom call backend config core ID fields.
191d12b5
Merge pull request #41162 from jabrailkhalil:feat/generalized-eigh
f716803b
[Mosaic TPU] Fix a data corruption issue.
3500ff42
[custom_jvp3] fix in-tree test failures under jax_custom_jvp3
2c687718
[checkify] check inside custom_jvp3 and custom_vjp3 primal functions
a03c5fd1
[hijax] read the custom_jvp3, custom_vjp3 and remat3 flags at call time
2b3f5c2b
Merge pull request #41187 from ROCm:users/gulsumgudukbay/drop-gfx1250…
5e71363f
Deprecate lax.psend and lax.precv.
f03d303d
[JAX:Pallas] Guard against Literal in `_jaxpr_has_unbatchable_dataflow`.
85e3c0b5
Use `call_sharding_rule` instead of `broadcasting_sharding_rule` sinc…
d57ed52e
Skip test cases that are failing in CI.
23ac35de
Merge pull request #41161 from mattjj:push-oqownpqmxslm
ba0b07eb
[mgpu] Push tcgen05 commit arrive predicate handling down into tcgen0…
4781218a
Merge pull request #41193 from mattjj:push-lurpzplqnpyw
d8d3ebcc
Reverts 191d12b55586e972f0b2ceebf5fd6ebfb38a4e4d
cdc97bca
Prevent a failing ordered IO callback from poisoning tokens forever.
ad7dc5b0
Merge pull request #41159 from mattjj:push-mkyxprtwtnxm
7236005c
[tests] Ignore psend and precv deprecation warnings in remaining GPU …
b47f1578
Merge pull request #41015 from bhaskargurram-ai:fix/pareto-support-mask
44ad65c2
[linalg] fix input shape validation in eigh & cholesky
81038a08
Add a wrapper to topologies to get a compile only device given a devi…
942b1407
Add ProfilerSession API and support for context based profiling
b4e35429
Sort lexsort keys in batches so compile time doesn't grow with the nu…
a63af0a5
[pallas:fuser] Simplify the scalar prefetch machinery
cfc62259
Fix vmap of a scan closing over a Ref batched on a non-leading axis
d1308a5a
Batch scan consts along their existing axis for all types
0f24a6d9
Batch a scanned-over Ref along its existing axis instead of moving it
51d1a5bd
Add Layout.for_array
d7dc3daa
[Mosaic] Split min reduction kind into minf/minsi/minui and similarly…
a424343b
Add layout rule for dot
aa85ffb1
[mgpu] Add predicate to tcgen05 commit arrive MLIR op.
34e4b2c6
Merge pull request #41210 from jakevdp:linalg-shape-error
8516670b
[mgpu] Add `predicate` to `tcgen05_commit_arrive`.
799aaa21
[Pallas/TPU] Emit DMA state effects for all underlying refs in pl.sel…
0e1a2714
Update XLA dependency to use revision http://github.com/openxla/xla/c…
9034ab8e
Fix typos in Blackwell matmul guide
fc530d96
Support type-2 and type-3 generalized Hermitian eigenproblems
c22ca60f
Fix expi at zero
797cb341
[typing] update for pyrefly v1.3.2
925018e4
[indexing] Treat fully replicated multi-device arrays as unsharded in…
fb6f015d
Merge pull request #41225 from matrix-kernel:docs-fix-typo
e5dbd72b
Merge pull request #41197 from jabrailkhalil:feat/scipy-eigh-problem-…
e0af9e17
Merge pull request #41221 from lindicaphxag-tech:fix-expi-zero
af92a4ae
Merge pull request #41186 from bhaskargurram-ai:perf/unique-axis
15998e80
[Pallas/Mosaic GPU] Add test demonstrating transposed `stmatrix` store.
7fc69a22
Relax reconstruction tolerance in QdwhTest.testQdwhOnRankDeficientInp…
8dc5907c
Merge pull request #39230 from wasim-builds:docs-conv-transpose
34ee5466
[JAX] Relax lax.cosh complex64 gradient tolerance on CPU and TPU.
c74c6d6c
Make the 4D dot layout test independent of the GPU GEMM backend.
600ad1b9
Merge pull request #41170 from Intel-tensorflow:minigoel/oneapi_ci_mu…
6075917d
Add google_benchmark microbenchmarks to jax/tests/numerics.
4a0b048d
SKip test_async_all_to_all_opt_barrier in OSS TPU until new libtpu is…
f36df970
Merge pull request #41234 from jakevdp:pyrefly-1-3
cbe2a56e
Merge pull request #41082 from Intel-tensorflow:bhavanis/update-docs
9d8ae3f1
Merge pull request #41191 from Intel-tensorflow:minigoel/oneapi_distr…
013726ad
[BugFix] Replace phase-2 PartialReduce with Pad-Reshape-Reduce in SPMD.
d9011abc
[hijax] make jax.vjp and HiPrim vjp rules interchangeable
d0571520
Add transpose layout rule
600095be
Merge pull request #41240 from mattjj:push-kzrwktmxlurw
83b55702
Write a better transpose layout test
a1603a6e
Integrate Triton up to 2074a1b8a5cd283ce584f8ebca825c4f0b215fc1
656227de
Update XLA dependency to use revision http://github.com/openxla/xla/c…
3c7bb6e3
Update XLA dependency to use revision http://github.com/openxla/xla/c…
d417226b
Update XLA dependency to use revision http://github.com/openxla/xla/c…
2346d416
Disable another overlap_test failing with libtpu < 0.0.50.
fbf6588d
Fix Poisson logpmf support for fractional loc
23efdd85
Update XLA dependency to use revision http://github.com/openxla/xla/c…
10d6b70f
Improve `jnp.hypot` numerical accuracy and add n-ary precision tests.
8bd835e3
[BugFix] Fix non-terminating expn branch loops under jnp.piecewise
afe04a2f
Merge pull request #39745 from adityasingh2400:fix-vmap-ref-in-axes
d5b6951f
[hijax] stay on the C++ jit dispatch path when only intermediates are…
35583fa7
Lower complex `abs` using `lax.hypot`.
22714b86
Merge pull request #41266 from mattjj:push-xkpqvsksmrzu
5bd8053f
Add accuracy tests for additional trig functions.
e6aeacf2
[pallas:mgpu] Allocation of semaphore references .
6b0a1afb
[typing] Use Self annotation for unary jax.Array methods
799aa057
Update XLA dependency to use revision http://github.com/openxla/xla/c…
456c052c
[Mosaic GPU] Diagnose unsatisfiable `Relayout` and register `Equals` …
f03c641a
[Mosaic GPU] Test SMEM <-> registers transfer error diagnosis on `Vec…
4a3d35a7
Update XLA dependency to use revision http://github.com/openxla/xla/c…
e3ae2a48
[Mosaic GPU] Propagate HBM alignment for peer references in `to_remote`.
f7045b3d
[numerics] Move mpmath domain and pole handling into test reference f…
19418b6d
[thunky] Add `jax.experimental.thunky` for XLA:GPU ThunkExecutor prog…
71b75c48
Fix reduceat accumulator dtype
709a7d62
Relax all-gather instruction name regex in SparseCore offload test.
192870c3
Automated Code Change
a06fd581
[CI] Require TPU v7+ for tests/overlap_test.py:test_async_psum_scatte…
a2daf819
Update libtpu to v0.0.49 and regenerate requirements lock files.
c317f2d2
[Mosaic GPU] Add `strict` attribute to `LayoutCastOp`.
e1d8d587
Merge pull request #40987 from hawkinsp:libtpu
fda59833
[Pallas] In TPU/GPU Interpret Mode, raises exception when a race is d…
764920d0
Switch to new RBE cluster for ROCm Bazel tests
0482de0b
Move `program_order` to its own module and add transformation rules.
fe341980
[Pallas:TPU] Fix pltpu.bitcast batching rule when batching over minor…
ecfdc482
Allow flattening of `AccRef` and `KeyScalarBundle` during lowering.
bf00fb63
Merge pull request #41220 from lindicaphxag-tech:fix-poisson-logpmf-f…
d89c9aa6
Merge pull request #41247 from SrijanRoy123-github:fix-reduceat-accum…
003297eb
Run DCE on Hi-jaxprs during lowering to lojax. This is similar to how…
a34d997f
Test interesting input values in numerics tests
658058ff
Simplify the xs batch-axis computation in _scan_batching_rule
d182777d
Make `program_order_p` as HiJAX prim so that opt_barriers are inserte…
e59f7f16
[Mosaic GPU] Preserve memref layout offset in `get_cluster_ref`.
4435223f
[CI] Skip test_ffi_custom_call_thunk_direct_and_extracted_from_jax.
99fadc28
Improve precision of `asin` fallback expansions.
0db49fff
Allow flattening of `AccRef` and `KeyScalarBundle` during lowering.
279f7b46
[Pallas MGPU] In GPU interpret mode, make barrier dealloc check lazy
3b3b5273
Merge pull request #41253 from athrva98:fix-expn-piecewise-hang
0da15eef
Merge pull request #41290 from mattjj:scan-batching-xs-axes-golf
db7ef515
Reimplement `lax.erfc` and `jax.scipy.special.erfcx` in JAX, with acc…
dc2d978c
[ROCm] Enable thunky on ROCm
5dca6dd6
[autodiff] avoid overflow in asinh jvp rule
2ec6ec86
[pallas:fuser] Trace index maps on placeholder scalars when comparing…
6882cc40
[pallas:fuser] Support Ref operands in Fusible.vjp_bwd
9553f8b6
[Pallas] Lower negative padding as a slice
3c8d620e
[lax] Removed an opt-in Pallas Triton lowering for ragged dot
b7914c2d
[pltpu.emit_pipeline] Add `output_wait_read_only` to emit_pipeline on…
434130eb
Default float16 to non-FTZ in numerics tests and fix subnormal underf…
ed4c61cd
[pallas:fuser] Keep the bwd's logs out of its cotangents when physica…
06310202
Reverts dc2d978c286467e817ed6c3149dbc9beb9024aba
70772107
Catch ValueError instead of IndexError for tuple.index in pure_callba…
be60eacb
fix tests
205c669c
Switch to rocm_library from rules_ml_toolchain
6c7cb1b1
Update XLA dependency to use revision http://github.com/openxla/xla/c…
ee97a8d3
Merge pull request #41309 from Cjkkkk:flex_attn
a954e248
Merge pull request #41310 from ROCm:switch_to_newest_rules_ml_toolcha…
08ebdf88
[jax] Reassociate constants in `jnp.linspace` and normal initializers…
eef7e60b
[numerics] Allow 1.0 ULP error for `float64` `sin` on ARM CPU in `tri…
cfa0d8dc
[thunky] Skip `thunky` tests on non-GPU backends and CUDA devices bel…
c72e0609
fix(scipy.stats): preserve precision in norm.isf for small tail proba…
0dc32a4c
Merge pull request #41305 from ROCm:users/gulsumgudukbay/thunky-rocm
7f5310fe
Relax query_grad/key_grad atol in `DotProductAttentionF8Test.test_sdp…
fdff3279
Add tests of numerical accuracy for gradients of `tanh` and `atanh`
598ebe55
[CI] Restrict `google-benchmark` to `python_version<"3.14"` in `build…
d2357e63
[Mosaic GPU] Narrow `collective` type in `LaunchContext.async_copy` a…
922602f3
[typing] fix jax.Array.astype annotation
06667818
Replace `PyTuple_SET_ITEM` and `PyList_SET_ITEM` with nanobind builders.
df3e275a
Merge pull request #41233 from ROCm:cph-upstream-use-new-rbe-cluster
3d87ca7b
Merge pull request #41324 from jakevdp:astype-annotation
fc8712e0
[Mosaic TPU] Extend support of reshapes of form (*src_major, src_mino…
40fc7a89
Merge pull request #40788 from mayuriphad:fix-callback-error-msg
c24f87b6
Merge pull request #41276 from Kuldeeep18:fix/norm-isf-precision
de856c17
Reimplement `jax.scipy.special.erfcx` and add `accuracy=HIGHEST` for …
e9b14a6f
fix instances of "the the" in JAX
1584d6dd
Merge pull request #41302 from P-bibs:lazy-barrier-dealloc-check
8ed533cf
[Pallas SC] Support remote transfers between HBM and vector subcore V…
b3b7d2e3
[pallas] Fix emit_pipeline under jax_enable_x64
1e0af75c
Update XLA dependency to use revision http://github.com/openxla/xla/c…
bb89a781
[mgpu] Add barrier_arrive_and_wait primitive.
4154b4a0
[Mosaic GPU][NFC] Pre-factor out `before_thread_sync` helper.
c12ad12f
[Mosaic GPU] Delete `WaitOp` from the dialect.
b0903459
Improve float32 log1p accuracy on TPU v5p, v6e and 7x.
11fa76a9
Improve erf_inv accuracy near |x| = 1.
4bbb23df
[Pallas:SC] Add dim to pltpu.bitcast and delete plsc.bitcast_p
7a1535cc
Change `ragged_dot` `group_offset` to a scalar array.
99dd9e4c
[Mosaic GPU][NFC] Raise error message in units of bytes instead of bits.
c5df14c8
[pallas/fuser] Improve custom_fusion validation and error messages ba…
516d99a5
[JAX - Pallas] Skip `test_pltpu_bitcast_minormost` when `libtpu < 0.0…
8eb6aff6
Update libtpu to 0.0.50 and regenerate lock files.
f4bb5a54
Preserve Jaxpr and JaxprEqn identity when DCE is a no-op.
b362f1cb
Fix single-core TPU failures and remove skips in multihost_utils_test
be4aa105
Merge pull request #40542 from Intel-tensorflow:kranipa/oneapi_linalg…
d0fe4ba4
[JAX - CI] Set `@ml_build_config_platform//:platform` in `ci_rbe_cache`.
d6763bee
Raise a descriptive ValueError when cache key __eq__ fails in weakref…
5c987ea9
[Mosaic GPU][NFC] Move comment related to synchronization around `tcg…
ad1a9927
[internal] internal type annotation fix
7ec1eb23
[pallas:mosaic] Added a nonreg test for `lax.cond` lowering
15e5e1c9
[Pallas/Mosaic GPU] Unconditionally run `_handle_transforms` on `Abst…
cafc741a
[mosaic_gpu] Added `async_copy_im2col` for TMA GMEM->SMEM copies in i…
21781983
[hijax] batch custom_jvp3/custom_vjp3 applications underneath, not vi…
b748e19c
Merge pull request #41223 from mattjj:push-ukyktnovpmqv
82e93582
[Mosaic GPU] Support partitioned and replicated TMA loads for cluster…
25be6532
Merge pull request #41355 from hawkinsp:libtpu
7e27a267
Minor simplification to `group_starts` calculation in ragged dot impl…
0f1b83b4
Use unstable sort in unweighted quantile calculations.
52bc9bf7
Update XLA dependency to use revision http://github.com/openxla/xla/c…
8a09c826
[Pallas/Mosaic GPU] Do not arrive on barrier in `copy_gmem_to_smem` w…
59fd3b22
[Pallas/Mosaic GPU] Fix `copy_gmem_to_smem` lowering when the barrier…
82d81ef0
[Pallas/Mosaic GPU] Fix `async_store_smem` when the barrier orders th…
6adbdcc7
Update XLA dependency to use revision http://github.com/openxla/xla/c…
b8af921d
[pallas:mosaic] Mark `PallasCallTest.test_cond_extui` as `@jtu.thread…
33072b35
[Pallas:TPU] Support pltpu.trace_value in TPU interpret mode.
69d9e172
[Mosaic GPU] Add `alloc_smem` and `alloca_scope` dialect operations.
c7df1f89
DCE per-branch jaxprs before merging residuals in `_cond_partial_eval…
ca2eaca1
[JAX] Avoid pre-broadcasting coefficients in `lax.polynomial` scan fa…
fb49a370
[mosaic_gpu] Added support for arbitrary window strides to `async_cop…
819cf0de
[typing] add internal `jax.random.Key` typevar
71475938
Preserve concrete `int32` loop bounds in Pallas Mosaic `emit_pipeline`.
378f9261
[ROCm] Run export_harnesses_multi_platform_test_gpu downstream
ea19b471
[ROCm] Stop running the continuous wheel tests on a schedule
967840be
[ROCm] Switch the pinned ROCm CI legs from 10.0 to 10.1
88190718
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Login to write a write a comment.
Login via GitHub