CI: 12/19/24 upstream sync #187
charleshofer
merged
154 commits into
rocm-main
from
ci-upstream-sync-66_1
Limit self-hosted jobs to JAX main repo
ae066c29
Merge branch 'jax-ml:main' into main
a67eae10
temporarily un-deprecate several jax.core APIs.
d3406768
Remove deprecated XLA GPU flags.
1021603f
Add flag desc to gpu_performance_tips.md
c563e473
Allow tracing and lowering (with lowering_platforms specified) to wor…
d0f63da4
Remove autotune sharing.
a123d4e3
Use a broadcasted gather in the sort JVP, rather than forming explici…
0922feb2
Merge pull request #25454 from sfvaroglu:sevin/slop_factor_doc
69263563
Merge pull request #25459 from hawkinsp:sort
5a3fa500
get_githash: fix support for missing git
1afed917
Migrate JAX MLIR Python dialect extensions to nanobind.
64eae324
Set device_assignment to None when only AbstractMesh exist in the com…
80cf1418
Merge pull request #25451 from jakevdp:undep-core
87b66f3c
Support FP8 for dot_product_attention
c67b6513
Merge pull request #25444 from justinjfu:gpu_docs_update
99b390ce
Merge branch 'jax-ml:main' into main
4de58e1a
Merge pull request #24536 from MichaelHudgins:main
a9c140d8
Update XLA dependency to use revision
078c7e44
Implement `process_call` and `process_custom_jvp` for LinearizeTrace
dea51cb7
Finalize deprecation of jnp.round_
c73f3060
fix `gamma_p` in vmap-based impl rule mode
b8304963
Merge pull request #25481 from jax-ml:custom-linearize-process-call
f4e5f14a
Add an experimental Cloud TPU presubmit job
d05ab5bb
Merge pull request #25487 from froystig:gamma-impl-vmap-fix
38592981
[Mosaic GPU] Fix layout inference traversal to traverse ops recursively.
23868383
Log DeprecationWarnings once per method/class
21d1d44e
Update XLA dependency to use revision
ab031577
[Mosaic:TPU] Replicating retilings with increasing tile size for (a) …
5493fabb
Update XLA dependency to use revision
c7d1c3d2
[export] Expand exporting to work with AbstractMesh.
afcb62ea
[mosaic_gpu] Allow calling `reduce_sum` on a fragmented array in spla…
ee7226d5
Drop the frequency of Cloud TPU tests
7e6877b2
Ignore a false-positive Hypothesis health check
16c44e51
Bump the oldest supported libtpu to conform to the 12 week window
e655358a
Merge pull request #22670 from wenscarl:fp8_sdpa
73cc7a57
Add Python.h include to fix Windows build.
11e0fdf3
Avoid assuming that jnp.sin will be traced in abstract mesh tests
3b9a8f79
Merge pull request #25435 from gnecula:export_abs
2b06f93c
Fix some flaky LAX autodiff tests.
a2e96995
Raise the timeout for Cloud TPU nightly CI
b3177daa
Relax tolerance for LAX reduction test in float16.
ed4e9823
Minor typo fixin doc.
cd7109e6
Fix incorrect capitalization in scan error message
74e9275b
Merge pull request #25511 from jakevdp:fix-scan-err
417d7a25
Update XLA dependency to use revision
302d8031
Add utility script and env for running the CI scripts under Docker
85528526
Re-deprecate a number of symbols from jax.core
cfa38841
Merge pull request #25508 from jakevdp:dep-core
772339ec
Fix `save_from_both_policies` in presence of `save_and_offload_only_t…
7dd401cb
[Pallas TPU] Use vector.broadcast instead of vector.BroadcastOp to fi…
72cb6215
Allow `lax.ragged_all_to_all` input and output operands to have diffe…
6d82a6fc
[Mosaic:TPU][NFC] Clean up local variable
7e6c52dc
[Mosaic:TPU][NFC] use_implicit_shape instead of Reshape + don't creat…
c6e6b11e
[Mosaic GPU] Add layout inference and initial lowering for `vector.{l…
03612554
[Mosaic TPU] Add support for bf16 abs
0ec902df
Make gmm TPU kernel tests significantly cheaper
3fc23712
Reduce the max_examples for splash attention tests
e1f037bf
Put abstract_mesh on every eqn so that we can preserve it during `eva…
473e2bf5
Cleanup: toward merging core.concrete_aval & xla.abstractify
2c722d9b
[Mosaic GPU] Add end-to-end lowering example for a pointwise kernel u…
36b12d58
[Mosaic TPU] Add support for the interleaved pack format to tpu.unpac…
4911a396
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
7de9eb20
Merge pull request #25473 from adamjstewart:build/no-git
7bbba40c
Update XLA dependency to use revision
7fe25791
Merge pull request #25456 from jakevdp:xla-abstractify
0fa54197
Streamline some core.concrete_aval compute paths
a2ac234f
Raise rather than return error
2518c623
Release v0.4.38
c4e3f5a3
Merge pull request #25538 from jakevdp:raise-runtimeerror
4aebe55f
Enable PJRT compatibility in cloud TPU CI.
71b23ea2
Merge pull request #25534 from jakevdp:faster-avals
05ad3934
Always use the same code for array avals
c9afc89c
Merge pull request #25544 from jakevdp:array-aval
d4031e99
Merge branch 'release/0.4.38'
fd8de7d2
Update version numbers after release.
ff52aedf
[Mosaic GPU] Skip testing uint64 unless 64-bit types are enabled
6ea47082
Add Pallas Philox implementation.
7e96914e
Merge pull request #25550 from hawkinsp:postrelease
8f4e13f6
Don't create a mesh inside jit in sharded-computation tutorial
83ca7c33
Merge pull request #25548 from andportnoy:aportnoy/mosaic-gpu-scalar-…
cce40666
fix squeeze lowering required sharding argument
0b49231e
Delete enable_memories code in C++ since that flag is always True and…
cca9afa2
Fix the breakage caused by deleted enable_memories config
1e221494
Remove core.concrete_aval and replace with abstractify
3cecbf34
Merge pull request #25555 from jakevdp:combine-avals
b56dc631
Allow `P.UNCONSTRAINED` in out_shardings at top level jit. This is re…
e854f165
[sharding_in_types] Add tests allowing inputs and outputs of jit to h…
09fdd0da
improve checkpoint / remat concreteness error with static_argnums
9acd4a95
[Mosaic:TPU] Fix bug after cl/707025084
f9737b95
Merge pull request #24516 from mattjj:improve-concreteness-error-in-r…
32627702
ref errors
42ac4ca3
[Mosaic:TPU] Allow null parts for tpu.pack_subelements, meaning "don'…
dc0b7747
Merge pull request #25449 from mattjj:ref-errors-3
96d4a758
Reverts b56dc63160eaccd7df05d03b1c38f804ff85f564
25524abc
Bump xla_extension_version after jaxlib release.
7776982a
Increase the minimum NumPy version to v1.25.
ee457184
Use StableHLO acos and update complex acos accuracy tests.
f592173c
Merge pull request #25559 from hawkinsp:minver
5c9756bb
[Mosaic TPU] Add support for exp, exp2 and log in bf16 on TPUv6
d95b95b4
Increase the timeout in profiler tests, but use threading.Event to ex…
7c1bea6b
Merge pull request #23830 from pearu:pearu/acos
5beb4794
[mosaic_gpu] Error on static out of bounds indices in `utils.parse_in…
98067fc1
[pallas:mosaic_gpu] Added a lowering rule for the general `lax.while_…
5e9cfce2
Use the right health check suppression
3c52ca6e
Merge pull request #25412 from shangz-ai:patch-3
66c8a2e9
Migrate StableHLO Python extension to nanobind.
3d54d035
[Mosaic GPU] Clean up imports in `gpu_dialect_test.py`.
6a03ea3e
Enable --config=clang only on newer Clang versions
bcca77cd
[Mosaic TPU] Add support for bf16 second minor reductions in TPUv6
6edfe9ea
Use capture_stderr instead of packing sys.stderr
2259a13d
Partial discharge for scan_p ops.
aaabb975
Migrate mhlo dialect extension to nanobind.
3f24dfd2
More linearize fixes
2be9a699
Merge pull request #25569 from hawkinsp:numpyver
464e5a27
JEP: effort-based versioning (EffVer)
88c89419
[JAX] Fix a small bug if shardings is tuple.
72e5ca95
Merge pull request #25490 from jax-ml:more-linearize-fixes
8abb1a73
Remove jax_enable_memories config flag. It defaulted to True for a ve…
8b734808
Update XLA dependency to use revision
53bff861
fix bug with jax.remat static_argnums not supporting int
bb16d5ac
[Mosaic TPU] Support direct cast i8 vector to mask
bf692efb
[Mosaic GPU][NFC] Split LaunchContext into a separate file
1ed5fa6b
[Pallas] Add version guard for non-32-bit selection in test and fix g…
74eca134
[Mosaic:TPU][NFC] Delete unused functions
13e721a2
[sharding_in_types] Check out_avals with mesh context too. This is be…
af63e443
[Mosaic] Verify that the target IDs are provided in remote DMAs
de359f5c
Re-land changes from https://github.com/jax-ml/jax/pull/25555
89a54a9e
remove whitespace from distributed arrays doc
c6576827
[Mosaic] Improve error verbosity of tpu.memref_slice verification
6bcec910
[Mosaic TPU] Support i32 vector multi reduction except cross lane.
3a5c4da4
add mutable array ref error checks to scan
e5285626
Merge pull request #25596 from mattjj:whitespace
46b18d27
Account for tokens in `allow_spmd_sharding_propagation_to_parameters`…
9041b02d
Merge pull request #25593 from mattjj:ref-errors-4
f65ecedd
[Pallas] Fix lowering tests for reduction ops
613a0cde
Refactor: move shaped_abstractify to core
676070f4
Add lax.composite primitive
d206cc3b
[Mosaic GPU][NFC] Move up the definition of `ThreadSemantics`.
16f0f97f
[Mosaic:TPU][NFC] Small cleanup of extui rule in apply-vector-layout
8188c574
[Mosaic GPU] Replace the dialect's layout enum with layouts holding t…
66ad2082
[Mosaic GPU] Add a new tiled layout, optimized for upcasting before W…
006c65d8
[Pallas:TPU] Use self.pallas_call to properly handle interpret mode
45159494
[Mosaic TPU] Guard tests for new features by the libtpu version
ad00ec1d
Merge pull request #25595 from jakevdp:mv-shaped-abstractify
76805325
jax.nn.one_hot: deprecate non-integer inputs
8c3c441e
[pallas:triton] Fixed a typo in a type annotation
af7a31f1
[jaxlib] Added a missing pytype_dep to :_triton_ext
c4fae4a7
[Mosaic GPU] Commit to using `Vector`s everywhere (and no `Tensor`s).
3915f4a1
Always suppress the differing_executors Hypothesis health check
23000a38
Merge pull request #25590 from jakevdp:fix-one-hot-float
8402a988
Merge pull request #25516 from jakevdp:jep-effver
49c92460
Merge branch 'rocm-main' into ci-upstream-sync-66_1
ab8937fa
Run CPU CI again
b8bbb148
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Login to write a write a comment.
Login via GitHub