ROCm/jax

Pull Requests Commits

fix-forward for pallas tpu memory spaces test

sharadmv committed 1 year ago

382b3e02

Remove unused internal optimization_barrier alias

Jake VanderPlas committed 1 year ago

d840447a

Make the params of more jaxpr primitives hashable.

hawkinsp committed 1 year ago

e04cc283

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

hawkinsp committed 1 year ago

b72be575

Use a frozenset for unconstrained_dims in sharding_constraint_p.

hawkinsp committed 1 year ago

688c3d33

Merge pull request #29441 from jakevdp:bincount-doc

Google-ML-Automation committed 1 year ago

109519a0

Add colorama back into test-requirements

jakeharmon8 committed 1 year ago

abb756d4

[doc] fix some inaccuracies in jnp.bincount docs

jakevdp committed 1 year ago

efbedc61

[Mosaic GPU] Enable transpose tests in mosaic_gpu.

dimitar-asenov committed 1 year ago

094b66fa

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

Google-ML-Automation committed 1 year ago

a9fdb768

Add `all_gather_invariant` to lax.

yashk2810 committed 1 year ago

6e2977d8

Reland the C++ safe_zip implementation.

hawkinsp committed 1 year ago

07193467

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

Google-ML-Automation committed 1 year ago

1228053f

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

sharadmv committed 1 year ago

f81f2589

Merge pull request #29425 from DanisNone:miss-dtypes

Google-ML-Automation committed 1 year ago

46b9eadd

Add hermetic `nvshmem` dependencies to JAX targets.

Google-ML-Automation committed 1 year ago

294d86b2

Temporarily disable AVX512 in linalg_test_cpu.

hawkinsp committed 1 year ago

4e3bf29b

add missing dtypes to jax.numpy.__init__.pyi

DanisNone committed 1 year ago

ef106037

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

bchetioui committed 1 year ago

83c292bd

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

hawkinsp committed 1 year ago

a2880363

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

4ad4d4a8

Merge pull request #29362 from gnecula:platform_index_linearize

Google-ML-Automation committed 1 year ago

1886a9e1

Merge pull request #29423 from rdyro:gpu_paged_fix

Google-ML-Automation committed 1 year ago

2c467d6a

Fix GPU quantized paged attention tests for < sm89

rdyro committed 1 year ago

e7d252dc

Add nightly linux jax wheel tests for python 3.14.0b1

kanglant committed 1 year ago

45e61d8c

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

yashk2810 committed 1 year ago

cc976cb0

Add a pytype disable around zstandard.

hawkinsp committed 1 year ago

3ef4db4d

Merge pull request #29129 from Matt-Hurd:rename_to_xprof

Google-ML-Automation committed 1 year ago

d00b1cad

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

Google-ML-Automation committed 1 year ago

004b5480

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

Matt-Hurd committed 1 year ago

97ab9e0b

Older