ROCm/jax

Pull Requests Commits

Deprecate 2d arguments to jax.numpy.cross.

hawkinsp committed 4 days ago

65d58c8a

Update ROCm runner labels to new labels

psanal35 committed 4 days ago

7214e7b8

ci(rocm): conv-zero A2-run -- serialize autotuner (XLA_ROCM_AUTOTUNE_SERIALIZE)

magaonka-amd committed 5 days ago

ec2f4baa

ci(rocm): conv-zero A1-run -- enable dedicated autotuner pool, immediate free

magaonka-amd committed 7 days ago

86837333

ci(rocm): conv-zero A-run -- enable GPU UAF guard (XLA_ROCM_FREE_GUARD), drop poison

magaonka-amd committed 7 days ago

942ee06e

ci(rocm): fix python matrix -- drop 3.11 (unsupported on jax-ml/main)

magaonka-amd committed 7 days ago

e67979f5

[ROCm] Re-enable testConvGeneralDilated on ROCm (conv-zero repro)

magaonka-amd committed 7 days ago

1122005e

ci(rocm): conv-zero continuous -- 4gpu x py3.11-3.14, build vs ROCm/xla@devtest

magaonka-amd committed 7 days ago

58988911

ci(rocm): conv-zero A-run -- expose BFC reuse race + poison + trace (4gpu)

magaonka-amd committed 7 days ago

70f55487

ci(rocm): revert pytest allocator to default (address -> silent CPU fallback)

magaonka-amd committed 7 days ago

8038ac75

ci(rocm): build against ROCm/xla@magaonka/devtest and use address allocator

magaonka-amd committed 7 days ago

3fe77bcd

[ROCm] Allow XLA_PYTHON_CLIENT_ALLOCATOR=address

magaonka-amd committed 7 days ago

081d7d30

Handle singleton dimensions in reshape sharding rule in a more robust way.

yashk2810 committed 8 days ago

e441101c

[Mosaic GPU] Support swizzle in `is_valid_assignment`

golechwierowicz committed 8 days ago

d1d1cd25

Update XLA dependency to use revision http://github.com/openxla/xla/commit/500d7b12672ca2377aa84f42fad18a7b6f89e2aa

Google-ML-Automation committed 8 days ago

786079b7

[trace_to_jaxpr] Port ops and TPU Pallas tests to use trace_to_jaxpr

gnecula committed 8 days ago

b040143f

Make sure cosmetic changes to async collectives code

yashk2810 committed 8 days ago

fbbc7884

[Pallas/Mosaic GPU] Delete debug artefact.

bchetioui committed 8 days ago

0d88e2d3

[Mosaic GPU] Correctly propagate `approx=True` to `FragmentedArray.ex2` in `FragmentedArray.exp(approx=True)`.

bchetioui committed 8 days ago

2fd83c70

[Pallas/Mosaic GPU] Add `__eq__` and `__hash__` implementations for `AbstractRefUnion`.

bchetioui committed 9 days ago

113cb4d2

[trace_to_jaxpr] Refactor mosaic/core.py for trace_to_jaxpr.

gnecula committed 9 days ago

e220bc25

[NFC] Use modern Python syntax for type variables.

allanrenucci committed 9 days ago

c24ca37d

Drop support for NumPy 2.0 and SciPy 1.14 for the jax 0.11.0 release.

hawkinsp committed 9 days ago

fcfd49da

Drop the deprecated public API module jax.cloud_tpu_init.

hawkinsp committed 9 days ago

dabbaba9

Remove stale JAX_IFRT_VERSION ifdefs.

hawkinsp committed 9 days ago

59eea209

Remove stale version tests now the minimum jaxlib is v0.10.2.

hawkinsp committed 9 days ago

86a46a6e

[Mosaic:GPU] Remove NVSHMEM usage from Mosaic XLA custom kernel and tests.

PatriosTheGreat committed 9 days ago

181661c3

Update GPU Interpret Mode to support Blackwell matmul0 example.

jburnim committed 9 days ago

2e4d7f65

[JAX] Update build.py on how to update the requirements files.

hawkinsp committed 9 days ago

fd078003

[trace_to_jaxpr] Refactor TPU Pallas call registration

gnecula committed 9 days ago

c8e6f6d0

Older