CI: 02/10/25 upstream sync #227
github-actions
merged
116 commits into
rocm-main
from
ci-upstream-sync-110_1
feat(gh-13291): Add exponential distribution functions: cdf, logcdf, …
42b64fc0
docs: Update `jax.scipy.rst` to include exponential distribution func…
46c5b865
chore: optimize calculations in exponential distribution functions fo…
a3f73073
fix: correct argument name in `promote_args_inexact` for `logcdf` fun…
f9c5b0d2
fix: update `logcdf` implementation to use `sf`
9a324bfb
[better_errors] Add debug info to more Jaxprs and WrappedFun (step 1)
d12aead6
Set cron to 05:00 UTC in tsan.yaml
a026de23
Readded missing part to use locally built numpy in the build in tsan.yml
fc408d20
Remove the FFI examples workflow cache.
42b8eb6c
[Mosaic GPU] Add simple Blackwell matmul example
aff2cba8
Use uv for dependency resolution on CI.
5db5e0d5
Merge pull request #26300 from vfdev-5:patch-2
2b944442
Merge pull request #26302 from vfdev-5:readd-missed-sed-part-in-tsan
09ee37a4
jax.lax: improve docs for exp & log functions
f180353d
Merge pull request #26307 from dfm:kill-ffi-cache
b1b88a36
fix: simplify `logcdf` implementation by removing unnecessary argumen…
8561f90f
Merge pull request #26078 from gnecula:debug_info_jaxpr
414449e1
Add get/set methods to TransformedRef
782215b0
Merge pull request #26310 from dfm:use-uv-in-ci
694572bc
![dependabot[bot]](https://avatars.githubusercontent.com/in/29110?v=4&s=60)
Bump actions/setup-python from 5.3.0 to 5.4.0
333c4e7a
Merge pull request #26278 from jax-ml:dependabot/github_actions/actio…
023f06f9
Merge pull request #26281 from jakevdp:lax-docs
fdcc04c3
[Pallas TPU] Add helpers for writing collectives
02f45313
Set the mesh as manual during partial_eval_custom in shard_map so tha…
307006e1
skip tests with extra requirements
1ae02bc0
Remove libtpu-nightly dependency from jax[tpu].
b1a2c27a
Merge pull request #26325 from hawkinsp:tpu2
781172c2
Merge pull request #26324 from mattjj:skip-tests-with-extra-requirements
9aa98130
Skip broken tests at HEAD
c07b6b52
Merge pull request #26159 from andportnoy:aportnoy/mosaic-gpu-blackwe…
d6be2351
[Mosaic GPU] Infer whether A/B are row- or column-major from strides
1fbc4a15
[Mosaic TPU] Add optimized casts for bf16->s4 in TPUv6
e7a4f893
[Mosaic GPU] Refactor the Blackwell matmul example and make it runnable
b79ab01e
Update XLA dependency to use revision
3a1c63c5
[pallas:mgpu] Cast all indices to i32 during lowering.
eeace3ce
[Mosaic GPU] Add helpers for dealing with TMEM references + implement…
f4dab0cf
#sdy Fix the format for the external link to jax-shardy-guide colab.
ae74d3e5
Run Bazel CPU/CUDA RBE jobs on pushes to main/release branches
b5b913ac
Replace Python 3.12 with Python 3.13 and add Python 3.10 to the matrix
bc055569
Roll back multinomial change from https://github.com/jax-ml/jax/pull/…
e4dac395
[better_errors] Add debug info to the Jaxprs formed for AD
abcaec70
Return arrays from `ArrayImpl._check_and_rearrange`.
4ae7fcf3
[Mosaic GPU] Write a higher-level tcgen05.mma helper reusing WGMMA im…
6d6c8c2e
[pallas_mgpu] Fix emit_pipeline_with_wgmma test.
9f53dfae
Refactor JAX wheel build rules to control the wheel filename and main…
d424f5b5
Compute buffer aliasing on a per buffer basis.
da0827b7
Merge branch 'main' into scipy-expon
7fc605f7
doc: add note about f16 casting in jnp.mean
9b402ecd
Merge pull request #26313 from gnecula:debug_info_vjp
c46b0215
Merge pull request #26259 from Qazalbash:scipy-expon
1eda5e2e
Merge pull request #26339 from jakevdp:mean-doc
10363663
[JAX] Add a test verifying the behavior of module-level state accesse…
f43d2b68
If sharding is not None (that's passed to convert_element_type), only…
0fb278a0
Update tsan job to upload to resultstore
be0f3d8f
disable head dim 256 on bw now
553199e4
[better_errors] Continue adding debug info to Jaxprs (step 3)
904b7486
Merge pull request #26279 from MichaelHudgins:tsan-resultstore
2e808f28
Merge pull request #26348 from gnecula:debug_info_jaxpr_3
5d647ccf
[pallas:triton] Temporarily reverted to the lowering using Triton IR
efbb0afd
[Mosaic GPU] Take TMEM as a TMEMRef in tcgen05.mma, not as a raw address
026b6c97
Increase test shard_count for shape_poly_test on GPU
a678396f
tweak an error message on mutable arrays
96b6f77c
[Pallas TPU] Remove obsolete skip condition
afad924d
[Mosaic GPU] Relax TMEM stride constraints on dimensions of size 1
a61d8002
[Mosaic GPU] Add more SMEM buffers to avoid blocking for memory traffic
6524e67f
Merge pull request #26363 from mattjj:mutable-array-error-tweak
7025b7d1
doc: improve docs for jax.lax trig functions
2fb750e0
Merge pull request #26342 from jakevdp:lax-docs
f54ddb31
Update the sparse GPU kernels in jaxlib to use the FFI.
5e915d33
Fix invalid lowerings for ROCm in Pallas
10c2374f
Temporarily disable source mapper on win32
49fa1274
Merge pull request #26369 from justinjfu:sourcemap_win
1b7b04f7
Start a new TPU interpret mode for Pallas.
1c82484c
Remove code for compatibility with jaxlib < 0.5.0.
840192d3
Add AMD ROCm GPU CI post-build check (#137)
ebf4a54f
Merge pull request #25097 from jburnim:jburnim_pallas_interpret_mode
4b86ff22
Default JAX_CPU_COLLECTIVES_IMPLEMENTATION to 'gloo'.
f07243a7
refactor: move scalar type defs out of lax_numpy.py
b4f98eef
DOC: update numpy/scipy versions in deprecation doc
e0ee4156
Merge pull request #26372 from jakevdp:split-lax-numpy
af29a7b0
Merge pull request #26375 from jakevdp:fix-doc
ce51384b
refactor: move array creation routines out of lax_numpy.py
7bacfbc6
Remove the tutorial about legacy custom calls.
21281f37
Merge pull request #26380 from dfm:remove-legacy-custom-call-doc
cce3df10
Update RNN kernels to use FFI.
c6e83903
Fix file path regexp for Windows
c1734f22
Update XLA dependency to use revision
4c402b7b
Add check to `lax.composite` to prevent `DynamicJaxprTracer` type err…
2828bce2
[Mosaic GPU] Add tests for various tcgen05.mma configurations
3b5e91b8
Merge pull request #26368 from ROCm:fix-rocm-pallas-lowerings
6ad38af4
[pallas:triton] Really revert to the lowering using Triton IR
35351f95
Merge pull request #26344 from Cjkkkk:disable_head_256_on_bw
e56c7dc5
refactor: move sorting ops out of lax_numpy
d3b3cd36
Remove the unused cu_cholesky_update kernel in favor of the FFI version.
5bc17f7e
[xla:python] Add a mechanism for "batch partitioning" of FFI calls.
c521bc62
jax.lax: better docs for hyperbolic trig functions
311e6683
Merge pull request #25810 from ROCm:gh-9948-add-gpu-ci-upstream
1ca8807d
Fix a missing "f" on an f-string.
d01520c6
DOC: make clear that printoptions are NumPy aliases
08563842
Merge pull request #26383 from jakevdp:jnp-sorting
3b470b95
Merge pull request #26409 from hawkinsp:fstring
7ffb613b
[mutable-arrays] persist shardings through xla computations
719031c1
[pallas:mosaic_gpu] Fixed a bug in how `delay_release` is handled in …
e5058079
Merge pull request #26376 from jakevdp:array-creation
ec477634
Merge pull request #26377 from mattjj:maintain-mutable-array-sharding
c0ba3626
Increase the absolute error tolerance to fix flaky tests.
cd075375
Speed up NamedSharding construction.
f21b0f03
Don't call `get_cur_mesh_sharding` if sharding-in-types mode is not e…
21e1be33
Merge pull request #26406 from jakevdp:lax-docs
a7c20e7f
Merge pull request #26407 from jakevdp:printoptions-doc
28903574
Update XLA dependency to use revision
b3df4d58
Update XLA dependency to use revision
8401d9b7
Use an PartitionSpec.UNCONSTRAINED to represent unconstrained dimensi…
cf308a84
Merge branch 'rocm-main' into ci-upstream-sync-110_1
31c1f254
Remove conflicting param for ci_build
4e403d29
Change to make CI run
9133253c
charleshofer
approved these changes
on 2025-02-11
Login to write a write a comment.
Login via GitHub