ROCm/jax

Pull Requests Commits

[ROCm test] Drop MIOPEN_FIND_ENFORCE workaround in pytest_rocm

Ruturaj4 committed 18 days ago

77b69680

Update XLA dependency to use revision http://github.com/openxla/xla/commit/256aa4e26041ae6d6b245203d19bb69854cc77b5

Google-ML-Automation committed 18 days ago

74ce6d42

Configure Bazel mirrors for JAX OSS builds.

Google-ML-Automation committed 18 days ago

7c3d6a1a

Create `pl.select_ref` to dynamically choose from multiple refs for DMAs.

IvyZX committed 18 days ago

70f3b181

Merge pull request #37362 from vfdev-5:jrand-api--rademacher-out_sharding

Google-ML-Automation committed 18 days ago

335ea36e

Fix and re-enable `test_remote_profiler_gcs_path`

Matt-Hurd committed 18 days ago

67dfed2c

Add a jaxpr_html IR dump mode that dumps a jaxpr to an HTML document.

hawkinsp committed 18 days ago

10ba2477

[Mosaic TPU] Support int16 abs.

yueshengys committed 18 days ago

82bc105e

[Mosaic] Refactor maskTiledVregs and selectWithBounds to consolidate usages of mask and select.

WindQAQ committed 19 days ago

25316030

Use --override_module to override XLA in Bzlmod builds

meteorcloudy committed 19 days ago

d8c90f0b

[Mosaic GPU] Support multidimensional transposes in layout inference.

golechwierowicz committed 19 days ago

1084c202

Bump cloud TPU version check to align with tpu_custom_call.py

WindQAQ committed 19 days ago

f9aa4709

[pallas] Added a new helper called `pl.align_to`

superbobry committed 19 days ago

2fb4349e

Merge pull request #37335 from Saumay:add-scipy-linalg-leslie

Google-ML-Automation committed 19 days ago

434b83d2

Roll back https://github.com/jax-ml/jax/pull/37117.

bchetioui committed 19 days ago

29bfc56c

[Mosaic GPU] Lower n-d transposes in WG semantics.

golechwierowicz committed 19 days ago

5eeaedb9

Added out_sharding to jax.random.rademacher sampler

vfdev-5 committed 19 days ago

f3321a4e

PR #37117: jax/mosaic: enable CUPTI V2 multi-subscriber coexistence on CUDA 13.2+

zhenying-liu committed 19 days ago

9056dd9e

Remove usages of xla_gpu_experimental_aot_compiled_thunks

beckerhe committed 19 days ago

23e47ff0

Update XLA dependency to use revision http://github.com/openxla/xla/commit/3629e8b9e828343902340c5c0eb3f288dd73c946

Google-ML-Automation committed 19 days ago

58e00d7b

Add float16 to f8 cast tests

apaszke committed 19 days ago

78ab0ab0

[Mosaic] Add bounds checks for dimension indices

Google-ML-Automation committed 20 days ago

c1225064

Update XLA dependency to use revision http://github.com/openxla/xla/commit/341f69b4e5505a52e444780c52d05b8790a39df1

Google-ML-Automation committed 20 days ago

a3ed1b2f

Update XLA dependency to use revision http://github.com/openxla/xla/commit/18261672d5c2704a6e1017a0b66bd70fabb7a445

Google-ML-Automation committed 21 days ago

2df8eced

Add jax.scipy.linalg.leslie

Saumay committed 21 days ago

b21f7fb3

Integrate LLVM at llvm/llvm-project@a7c97252ed18

Google-ML-Automation committed 21 days ago

189907cf

Add subprocess registration hooks to JAX profiler libraries.

bmass02 committed 21 days ago

d1f6d783

Add hijax support to freeze/empty_ref

sharadmv committed 21 days ago

a7370b57

[Pallas TPU] Add support for VectorSubcore + TensorCore async MPMD

sharadmv committed 21 days ago

da5e7bbc

Merge pull request #37321 from samanklesaria:orthogonal_out_sharding

Google-ML-Automation committed 21 days ago

8f2e8fb2

Older