ROCm/jax

Pull Requests Commits

[ROCm] Fix bool/string type error in blocking-gate workflow inputs

mminutoli committed 137 days ago

7406c5d8

[ROCm] Restore conditional buffer_callback_test_gpu exclusion for pre-built wheels

mminutoli committed 137 days ago

99fbfce6

[ROCm] Remove TheRock from blocking gate (ROCm 7.2.0 only)

mminutoli committed 137 days ago

c8503f47

[ROCm] Drop continue-on-error from blocking-gate (unsupported with uses:)

mminutoli committed 137 days ago

45220b1c

[ROCm] Pass halt-for-connection to blocking-gate on workflow_dispatch

mminutoli committed 137 days ago

606b991d

[ROCm] Rework blocking PR gate: file-driven targets, rocm_gate config, explicit target files

mminutoli committed 137 days ago

3ed2fa4b

[ROCm] Add blocking Bazel PR gate workflow

mminutoli committed 137 days ago

c0ecae1c

[ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images

gulsumgudukbay committed 138 days ago

b0ab8909

[ROCm] Embed TheRock-aware wheel RUNPATH at link time

gulsumgudukbay committed 138 days ago

d430e3bb

[ROCm] Add TheRock 7.13 and TheRock latest to CI test matrices

mminutoli committed 138 days ago

dfb1a195

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

yashk2810 committed 138 days ago

e441101c

[Mosaic GPU] Support swizzle in `is_valid_assignment`

golechwierowicz committed 138 days ago

d1d1cd25

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

Google-ML-Automation committed 138 days ago

786079b7

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

gnecula committed 138 days ago

b040143f

Make sure cosmetic changes to async collectives code

yashk2810 committed 138 days ago

fbbc7884

[Pallas/Mosaic GPU] Delete debug artefact.

bchetioui committed 138 days ago

0d88e2d3

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

bchetioui committed 138 days ago

2fd83c70

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

bchetioui committed 139 days ago

113cb4d2

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

gnecula committed 139 days ago

e220bc25

[NFC] Use modern Python syntax for type variables.

allanrenucci committed 139 days ago

c24ca37d

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

hawkinsp committed 139 days ago

fcfd49da

Drop the deprecated public API module jax.cloud_tpu_init.

hawkinsp committed 139 days ago

dabbaba9

Remove stale JAX_IFRT_VERSION ifdefs.

hawkinsp committed 139 days ago

59eea209

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

hawkinsp committed 139 days ago

86a46a6e

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

PatriosTheGreat committed 139 days ago

181661c3

Update GPU Interpret Mode to support Blackwell matmul0 example.

jburnim committed 139 days ago

2e4d7f65

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

hawkinsp committed 139 days ago

fd078003

[trace_to_jaxpr] Refactor TPU Pallas call registration

gnecula committed 139 days ago

c8e6f6d0

Merge pull request #38559 from hawkinsp:postrelease

Google-ML-Automation committed 139 days ago

6b7b8ac8

Update requirements lock files after JAX 0.10.2 release.

hawkinsp committed 139 days ago

31c559c7

Older