ROCm/jax

Pull Requests Commits

[ROCM] Fix export harness tests

Ruturaj4 committed 2 years ago

35c70fd3

[Mosaic TPU] Enable lane broadcast for packed types and offsets outside of first tile, and fix some broadcast infer logic

tlongeri committed 2 years ago

0f834cdf

Use vmap for random_gamma implementation on CPU backend

ezhulenev committed 2 years ago

15d43892

Make the vmap(jit) or vmap(wsc) with a concrete layout error more informative

yashk2810 committed 2 years ago

2eb1888c

Merge pull request #22673 from mattjj:cond-flakey-test

a-googler committed 2 years ago

6f68887e

deflake cond memory leak regression test

mattjj committed 2 years ago

f72a3f8e

Merge pull request #22674 from mattjj:remove-pdot

a-googler committed 2 years ago

f75e52db

[XLA:Mosaic] Support any int type upcast.

bythew3i committed 2 years ago

b1b7d046

Fix "unhashable type" error when passing a jax array as the "repeats" argument to jnp.repeat().

hawkinsp committed 2 years ago

1ac20854

remove pdot and xeinsum (since xmap is gone)

mattjj committed 2 years ago

88d1cd73

Delete mesh.Loop now that xmap has been deleted

yashk2810 committed 2 years ago

7de3c061

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

c4d3c8dd

Merge pull request #22669 from hawkinsp:repeat

a-googler committed 2 years ago

3ed9acba

Merge pull request #22668 from cool-RR:nanoseconds

a-googler committed 2 years ago

92806ee9

Added argument validation to mosaic_gpu_init_tma_desc

superbobry committed 2 years ago

5e418f5a

Simplify jaxpr for jnp.repeat in scalar repeat case.

hawkinsp committed 2 years ago

f07e963b

Merge pull request #22619 from jaro-sevcik:rename-mock-gpus

a-googler committed 2 years ago

d4e08a98

Merge pull request #22665 from rajasekharporeddy:testbranch1

a-googler committed 2 years ago

5d352a8b

Show elapsed time in nanoseconds

cool-RR committed 2 years ago

0d92d310

Better docs for jnp.triu_indices_from and tril_indices_from

rajasekharporeddy committed 2 years ago

d7171355

Merge pull request #22664 from jakevdp:astype-device

a-googler committed 2 years ago

f17d0f38

Merge pull request #22663 from jakevdp:array-api-cleanup

a-googler committed 2 years ago

593afa67

[array api] streamline astype device implementation

jakevdp committed 2 years ago

81b9db6b

register several deprecations in jax.numpy

Jake VanderPlas committed 2 years ago

31abce1c

Activate Cholesky Factorization Kernel to XLA's FFI

Paweł Paruzel committed 2 years ago

ae40c879

[array API] clean up some superseded definitions

jakevdp committed 2 years ago

8c36f90c

Merge pull request #22653 from rajasekharporeddy:testbranch1

a-googler committed 2 years ago

9ea79c61

[mosaic_gpu] arrive_expect_tx() also accepts index typed values

cperivol committed 2 years ago

d54bf529

Increase shard count for //third_party/py/jax/tests:lax_numpy_ufuncs_test_cpu.

belitskiy committed 2 years ago

ba50e774

Fixed Pallas Mosaic GPU tests

superbobry committed 2 years ago

99a8b92a

Older