CI: 11/06/24 upstream sync #131
charleshofer
merged
111 commits into
rocm-main
from
ci-upstream-sync-11_1
Re-land #24589 with fixes to handle `dtype` that is not compatible wi…
0181cb39
Remove the `variant` since sparsecore is only on v5p and it's device …
7f4a34e1
Remove UnshapedArray values from JAX (it remains as an abstract class).
f355dcf3
[Mosaic] Extend tpu matmulop to have dimension dims. Add support for …
5aeffde7
[MOSAIC:GPU] Extend the mosaic mlir dialect with fragmented layouts.
7d504cd9
[pallas:mosaic_gpu] `plgpu.copy_smem_to_gmem` no longer transparently…
85662f6d
[Mosaic GPU] Add Python bindings for the Mosaic GPU MLIR dialect.
c708a04c
Run dot algorithm tests with PJRT plugin.
52ad6052
Merge pull request #24480 from dfm:dot-algorithm-plugin-enable
8abedda8
Move the CUDA end-to-end example to FFI examples workflow + hosted
ce8dba98
[Mosaic] Fix some imports.
7d9f5656
Reverts 72f9a493589a1046e6927a5f16d7dc71df530743
a75d9462
Merge pull request #24190 from dfm:ffi-examples-gpu
7ff5a4ea
[Mosaic] Add extension files for infer/apply vector layout.
8296f6e0
Remove implicit sharding annotation for tpu custom call.
c758373b
Update XLA dependency to use revision
8536eca4
Remove ConcreteArray from JAX. It's easy to do trace-time concretizat…
48f24b6a
[Pallas:TPU] Use arith.divui for uint32 div.
7af7a60d
[array api] update test suite to latest commit
17ad8a95
Remove a number of expired deprecations.
2b9c73d1
Add a register_dataclass example to the pytree tutorial.
467bd09f
Simplified conditional in flash attention.
423cd2ad
Merge pull request #24647 from emilyfertig:emilyaf-doc-pytree-dataclass
5a3ed6c7
Add a JaxIrContext that subclasses mlir.ir.Context and avoids calling…
84c8794b
[Pallas TPU] Add lowering for `lax.nextafter`
f60b97ce
Merge pull request #24649 from jakevdp:array-api-update
f4d675ea
Merge pull request #24652 from jakevdp:old-deps
2a41c04f
remove busted example from shmap jep
26f70c9c
Merge pull request #24665 from mattjj:shmap-jep
24f3a2bb
Fix array API tests.
e657a4b2
Merge pull request #24667 from jakevdp:fix-array-api
bd7c3019
Fix #24661: Add zsh support to conda install documentation
5e1366c4
[Mosaic] Set TPU CustomCall device type based on the core_type attribute
f462d7e5
Update pre-commit workflow to cache on jax version
b7bdee90
Fix signatures test: new axis argument in trim_zeros
97e8a4c8
Merge pull request #24668 from Li-Jesse-Jiaze:fix-issue-#24661
453d6ff0
Merge pull request #24669 from jakevdp:fix-sig-test
7bc026e4
Set minimum supported Python version to 3.10 for matplotlib.
a0b0a8e5
Update XLA dependency to use revision
d606c242
[sharding_in_types] Allow `device_put` to reshard inputs. `device_put…
07858fa9
Add `compiler_options` argument to `jax.jit`.
fff33f90
Verify we can offload more than a single computation to SparseCore
d38da5d1
support exec_time_optimization_effort and memory_fitting_effort xla c…
0f3ba425
[export] Cleanup in the export module.
292a00b3
[jax] Mesh discharge rule should return None for inputs it did not to…
72eb5088
Update XLA dependency to use revision
d679c0ab
Remove lattice system from JAX, especially raise_to_shaped (except as…
ec39b592
Disable breaking testAutoPgleWithPersistentCache. The test is unstabl…
d4ddabb1
Fix ColocatedPythonTest. The test has been failing only on pytest nig…
e1270533
Update XLA dependency to use revision
2d94914c
[pallas:mosaic_gpu] Added a 2D test for `emit_pipeline`
c52b3227
Reverts ec39b592f7c096b0b8183723feaab2ed0d001041
f281c6f4
issue #24691
95146deb
Fix `_cuda_path` for case when `cuda_nvcc` is a namespace package
9df719f8
#sdy Fix Shardy bug where we weren't setting shmap in/out shardings a…
3544efca
[pallas:mosaic_gpu] `lax.fori_loop` lowering now promotes the carry t…
d2bbd564
[tpu] Disable a cumulative reduction test on TPU v6e that currently h…
f9541700
Merge pull request #24684 from hartikainen:fix-cuda_path
15febbf0
Merge pull request #24692 from jiaxi98:cholesky_document
26c0c5c7
Remove the `initial` argument to `jax.nn.softmax` and `jax.nn.log_sof…
e9acaa84
Merge pull request #24543 from dfm:pre-commit-cache-key
e80ead51
Update XLA dependency to use revision
38b4d001
Update link to algebraic_simplifier.cc to point to OpenXLA instead of TF
74da736e
[JAX] [XLA:Python] Move JAX configuration objects into C++.
ab47d468
Update JAX landing page - Flax
9f0e6237
Update details on JAX libraries in JAX README.md
b976b1ab
Merge pull request #24703 from 8bitmp3:update-landing-flax
ea4df795
Merge pull request #24702 from 8bitmp3:patch-3
2313b86f
rollback due to data race
a913fbf2
Fix Typos
a80d027d
[Mosaic GPU] Ensure that the dialect module can be loaded successfully.
63e59c5f
[mosaic_gpu] Check the return code of `gpuEventCreate` and `gpuEventD…
34b4787e
Improve efficiency of jax.scipy.stats.rankdata
5f90f63d
Reverts f281c6f46475270a57a02416469226315377592c
478b750c
Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461
0e8acff5
[jax.distributed] Enable grpc channel compression
7bdb2bf9
Make Tracers non-hashable
095bb0e7
Adding `JAX_LOGGING_LEVEL` configuration option
04f2ef9e
Merge pull request #23969 from gspschmid:gschmid/dist-compression
a70e7993
Merge pull request #23468 from rdyro:rdyro-add-logging-env
497a5a35
Merge pull request #24710 from rajasekharporeddy:typos
c1af808c
Merge pull request #24717 from jakevdp:fix-rankdata
939b41f5
Update XLA dependency to use revision
21f4b085
Merge pull request #24704 from andportnoy:patch-3
563ecdf2
logging fixes
b60d0ab8
Include mpmath as a bazel dependency of lax_test.
ea1e8795
Merge pull request #24729 from rdyro:logging-fix
4b4fb9da
Add support for tpu v5e to `jax.make_mesh`
b8c263a5
Merge pull request #24722 from jakevdp:tracer-hash
0fcb2f39
scipy.special.beta: remove deprecated x and y parameters
d698da61
Fix a bug in jax.scipy.stats.rankdata leading to breakage with shape …
542cb2e5
Merge pull request #24723 from jakevdp:beta-dep
cbaafbbe
[Mosaic] Verify that tpu.sem_wait semaphore rank is zero
3df204a4
[Pallas:TPU] Fix some stale/wrong skip conditions.
b6f5c95a
charleshofer
approved these changes
on 2024-11-12
Fix debug_nans false positive in jnp.quantile
0720942b
Remove some obsolete deprecation registrations
f1caa0ed
Update XLA dependency to use revision
95f7b247
Adding start index and kv_seq_len to decode kernel
8463eb08
Remove some obsolete deprecation registrations
f1caa0ed
Update XLA dependency to use revision
95f7b247
Fix yaml checker
ec3f5006
Set runners for ROCM
8faf2311
Allow devs to kick off sync job manually (#119)
bf035083
Unpin container in CI build and remove libssl-dev install
909f746d
Rename the CI flow to 'ROCm CI' and only run it on PRs to rocm-main (…
d2bda084
Fix nightly sync permissions (#124)
249ce156
Add GHA workflow for opening PRs upstream (#116)
f09863c4
Create a new branch when merging upstream main to rocm-main (#128)
78310661
Fix upstream sync checkout (#130)
a5ee6dc3
Fix FFI example test in CI
144bef02
Merge branch 'rocm-main' into ci-upstream-sync-11_1
587d7331
Login to write a write a comment.
Login via GitHub