CI: 02/03/25 upstream sync #219
github-actions
merged
123 commits into
rocm-main
from
ci-upstream-sync-104_1
add segment packing
28b642aa
Added CI job with TSAN and free-threading
ce597af5
Add test case for jit compile of vmap on gpu.
1ff268e2
address comments
bf4c3a77
CI: simplify nightly workflow definition
36c6c74c
Integrate LLVM at llvm/llvm-project@2e5a5237daf8
e332b94f
Don't apply shardy config during def_partition.
55d891c5
[Mosaic GPU] Add support for tiled loads/stores with sub-byte types
f504d324
[Mosaic TPU] Optimize clipping impelmentation in arith.fptosi
29b658b3
[Mosaic GPU] Add specialized support for some int4 -> bfloat16 casts
a4fe5c1a
Update XLA dependency to use revision
ddc8beff
DOC: add discussion of exceptions in pure_callback
ba2858f8
Merge pull request #26131 from jakevdp:fix-nightly-ci
5625b2fd
Add win32 guard to fix imports on Windows
7a4a53ad
[sharding_in_types] Upstream changes from defaulting sharding_in_type…
8f248fe6
DOC: avoid deprecated argument in external callbacks
25aa5a30
[Mosaic GPU] Implement a lowering for the dialect WGMMA op
d9f67ffe
Simplify dtype inference in lax.linalg.eig abstract eval rule.
09392d81
Merge pull request #24898 from vfdev-5:add-tsan-ft-ci-job
95558a9f
[better_errors] Expand debug info testing with eager mode, and MLIR m…
f8673cde
Merge pull request #26157 from jakevdp:callbacks-vectorized
fa9bb231
disable kernel test due to races.
23c8607b
#sdy Enable test_partial_auto_of_random_keys under Shardy.
9cbff642
Add support for `axis_name` and `axis_index_groups` to `lax.ragged_al…
809e1133
Make `JAX_CPU_COLLECTIVES_IMPLEMENTATION` and `JAX_NUM_CPU_DEVICES` e…
2aa810fe
Merge pull request #26154 from jakevdp:pure-callback-doc
bf22b53c
Always dispatch CPU executables synchronously when they include callb…
83457c11
[mosaic_gpu] Allow tiled array instead of wgmma.
f2f7a150
Merge pull request #26151 from gnecula:debug_info_collect_lowered_jaxprs
a459e7e4
Disable async dispatch within the body of a host callback.
9d39ab30
Relax static offset restriction on memref_ptr for sub-byte types
c9dfdb4e
[export] Fixes for export_harnesses_multi_platform_test.
8720e955
Revert https://github.com/jax-ml/jax/pull/25982 since callbacks can n…
e2eff1f8
Update XLA dependency to use revision
c01273c6
Internal: avoid adding _DimExpr to dtypes._weak_types
955e7c47
[jax:custom_partitioning] Make propagate_user_sharding default to None.
20843643
[sharding_in_types] Add vmap + explicit sharding support. The main ch…
dcb28f12
Unconditionally skip async deadlock test for pure_callback.
2ae018ed
Update the list of tsan suppressions.
4f00d451
Merge pull request #25980 from codinglover222:jit-vmap-compile-test
0a30ef3c
Increasing shard count and removing asan builds to prevent timeouts.
8c02731a
Pulled `DLDeviceType` to XLA backend mapping into a global
d4ced960
Merge pull request #26185 from hawkinsp:tsan
a1e4121b
Merge pull request #26188 from dfm:revert-callback-docs
152099ee
improve runtime error message
186478d2
Add skeleton for a multi-pass source mapper for Jaxprs/HLO to jax.exp…
b01111d9
Allow multiple meshes for avals but in that case, just use empty_abst…
d223dfc3
[better_errors] Refactor more uses of pe.tracing_debug_info (part 3)
32c98b9a
Merge pull request #26100 from gnecula:debug_info_no_pe_debug_info_3
4c8d7379
[Mosaic GPU][NFC] Fix wrong type annotations, and do some NFC cleanups.
46512e68
[Mosaic GPU] Add ArriveExpect and Wait ops on dialect barriers with e…
6214c25a
Skip ragged collective tests on CPU.
19c17bb2
[XLA:CPU] Add CPU client support for layout modes.
9442f90c
Added cpython cache step
07a7b7de
Update XLA dependency to use revision
c3814921
Merge pull request #26189 from vfdev-5:add-ci-cache-for-cpython-build
4af0481b
Merge pull request #26150 from jreiffers:main
1003ba93
[Mosaic GPU] Eliminate the `arrive` attribute from `mosaic_gpu.async_…
d8f3b33a
Return arrays from `ArrayImpl._check_and_rearrange`.
bb951136
Add TSAN numpy step
f4199fc4
[Mosaic][NFC] Factor out xla-array related utils in a separate file.
d4758b6d
Merge pull request #26208 from dfm:disable-ragged-test
2e40549c
Try to match out_spec with in_spec if both shardings are full auto an…
f4e2c6c3
Fix `lax.ragged_all_to_all` degenerate case
a8df383c
Do automatic casting from auto -> manual when the context mesh is man…
9107ee4a
[array api] bump test suite to latest commit
7f4796a2
Disable source mapper test for optimized hlo
834e0d7c
address lint and typecheck
ba6b1fdd
remove checks since they are redundant and we can change out_aval bec…
1f33cad3
Merge pull request #26221 from justinjfu:remove_smap_test
9dfe03c5
[Mosaic TPU] Support non-32 bit mask relayout
785a63ad
Disable a debug_info_test test that fails in CI.
a2f7824c
Pass filter=data to tar extractall to avoid a warning under Python 3.12+
0705ec2c
Fix CI failures in xla_metadata_test.
60fad99c
Merge pull request #26218 from jakevdp:bump-array-api
efacec4c
[Mosaic GPU] Add support for tiled swizzle=16 (i.e. no swizzle) loads…
10ac6b7e
[Mosaic GPU] Allow uneven partitioning of dimensions into tiles in Ti…
cadfcc7a
[CI] Update naming comments in workflow files.
873db96f
Update XLA dependency to use revision
e6e7621f
Reverts 86643a1b3e0516e1a2ddbdabbb714cf8c0301f18
1bfdd504
Re-enable Pallas bf16 exp2 tests on TPU
c1e13605
[mgpu] Correct instruction for conversion of unsigned int types.
bf967173
Disable two tests of GC behavior if there are multiple threads per pr…
2d9cd86a
test: change random matrix generation for Rotation
4d433d06
jax.lax: improve documentation for several functions
522200ff
Fix GCS download URI in Bazel CUDA non-RBE workflow
04476cf1
Merge pull request #26220 from jakevdp:lax-docs-1
ccef8301
Adjust conf.py to remove prompt from sphinx copybutton
dd7c203b
Lower `np.ndarray` to `DenseElementsAttr` instead of `ArrayAttr`.
20555f63
Avoid call to asarray in jnp.einsum
4e30a08e
[pjrt] Use the `PjRtMemorySpace` version of `BufferFromHostLiteral`
5cad02c1
Merge pull request #26236 from adonath:sphinx-copybutton-no-prompt-copy
def5fd52
delete unnecessary line
0df7f182
Merge pull request #26194 from jax-ml:fix-dist-init-runtime-error
cb79ff4d
Merge pull request #26245 from jakevdp:fix-line
70aed64a
Merge pull request #26239 from jakevdp:rotation-test
b039f976
Merge pull request #26238 from hawkinsp:coretest
0ef2ccfd
Add additional tsan suppressions for races.
d4c66578
Merge pull request #26247 from hawkinsp:tsan2
0bf8462c
Fix dtype issue in stirling approximation
216bd9a6
Merge pull request #26195 from vfdev-5:use-tsan-numpy-in-tsan-ci
3f737588
Reverts bb951136e9b91a584bb422119ada76cc69c86024
3b2410f7
Add mode='fan_geo_avg' to nn.initializers.variance_scaling.
96d3447e
Merge pull request #26249 from jakevdp:fix-sterling
a9f4dd71
Merge pull request #25766 from carlosgmartin:nn_initializers_variance…
872e6c0e
Move `ragged_all_to_all` test under appropriate test file
c4e17632
[Mosaic TPU] Support dynamic gather along axis 0 or 1 for 32-bit vreg…
d8b92113
[Pallas TPU] Support jnp.take_along_axis for 32-bit vreg-sized vector.
ed952c8e
Update XLA dependency to use revision
eb04fcbe
[pallas] Support DMA start partial discharge and run_scoped() does it…
8649132d
[pallas/pallas_mgpu] Discharging run_scoped should not be discharging…
b23f8f41
Merge pull request #26174 from skye:cpu_configs
de48ce2a
Update XLA dependency to use revision
af84143e
Merge pull request #26243 from jakevdp:einsum-asarray
57fa3721
Reverts 3f737588fd49fcac36d6bd00cfdabae07df1fbb2
58d8a97f
Stop passing None into jax.NamedSharding in preparation for followup …
da97ee25
[pallas_mgpu] For loops can have **non-ref** accumulators for carries.
b48d15d7
Reject invalid None in jax.NamedSharding(spec=None).
cb188a0c
fix lint
8c4d6d69
Add FFI example demonstrating the use of XLA's FFI state.
28afd252
Update XLA dependency to use revision
49821e81
Merge pull request #25812 from Cjkkkk:segment_ids
7164c6ba
Merge branch 'rocm-main' into ci-upstream-sync-104_1
ec2c3050
charleshofer
approved these changes
on 2025-02-03
Login to write a write a comment.
Login via GitHub