ROCm/jax

Pull Requests Commits

fix-forward for pallas tpu memory spaces test

sharadmv committed 306 days ago

382b3e02

Remove unused internal optimization_barrier alias

Jake VanderPlas committed 306 days ago

d840447a

Make the params of more jaxpr primitives hashable.

hawkinsp committed 306 days ago

e04cc283

[jaxlib] Change Traceback to be a raw CPython class rather than a nanobind class.

hawkinsp committed 306 days ago

b72be575

Use a frozenset for unconstrained_dims in sharding_constraint_p.

hawkinsp committed 306 days ago

688c3d33

Merge pull request #29441 from jakevdp:bincount-doc

Google-ML-Automation committed 306 days ago

109519a0

Add colorama back into test-requirements

jakeharmon8 committed 306 days ago

abb756d4

[doc] fix some inaccuracies in jnp.bincount docs

jakevdp committed 306 days ago

efbedc61

[Mosaic GPU] Enable transpose tests in mosaic_gpu.

dimitar-asenov committed 306 days ago

094b66fa

[Pallas TPU] Small fix to memory space constraints on pallas_call inputs.

Google-ML-Automation committed 306 days ago

a9fdb768

Add `all_gather_invariant` to lax.

yashk2810 committed 306 days ago

6e2977d8

Reland the C++ safe_zip implementation.

hawkinsp committed 306 days ago

07193467

[JAX] Update the example to use jax.numpy rather than numpy.

Google-ML-Automation committed 306 days ago

1228053f

[Pallas TPU] Support memory space constraints on pallas_call inputs.

sharadmv committed 306 days ago

f81f2589

Merge pull request #29425 from DanisNone:miss-dtypes

Google-ML-Automation committed 306 days ago

46b9eadd

Add hermetic `nvshmem` dependencies to JAX targets.

Google-ML-Automation committed 306 days ago

294d86b2

Temporarily disable AVX512 in linalg_test_cpu.

hawkinsp committed 306 days ago

4e3bf29b

add missing dtypes to jax.numpy.__init__.pyi

DanisNone committed 306 days ago

ef106037

[Mosaic GPU] Add conversion logic for `i4 -> f8e4m3fn`.

bchetioui committed 306 days ago

83c292bd

Move NamedSharding.__eq__ and NamedSharding.__hash__ into C++.

hawkinsp committed 306 days ago

a2880363

Update XLA dependency to use revision

Google-ML-Automation committed 307 days ago

4ad4d4a8

Merge pull request #29362 from gnecula:platform_index_linearize

Google-ML-Automation committed 307 days ago

1886a9e1

Merge pull request #29423 from rdyro:gpu_paged_fix

Google-ML-Automation committed 307 days ago

2c467d6a

Fix GPU quantized paged attention tests for < sm89

rdyro committed 307 days ago

e7d252dc

Add nightly linux jax wheel tests for python 3.14.0b1

kanglant committed 307 days ago

45e61d8c

Add execution to unreduced tests now that it works end-to-end

yashk2810 committed 307 days ago

cc976cb0

Add a pytype disable around zstandard.

hawkinsp committed 307 days ago

3ef4db4d

Merge pull request #29129 from Matt-Hurd:rename_to_xprof

Google-ML-Automation committed 307 days ago

d00b1cad

Add alternative location of `CUDA_ROOT` for Bazel build/tests with hermetic CUDA.

Google-ML-Automation committed 307 days ago

004b5480

[XProf] Change tensorboard-plugin-profile to new xprof package

Matt-Hurd committed 307 days ago

97ab9e0b

Older