ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

40d569b2

document vmap peculiarity of experimental RNG implementations

froystig committed 1 year ago

f30ebd85

improve RNG doc around implementation configuration

froystig committed 1 year ago

6ddd488d

Merge pull request #22080 from vfdev-5:add-device-kwarg-linspace-array

a-googler committed 1 year ago

aeff5b61

Un-skip //third_party/py/jax/tests:pytorch_interoperability_test_cpu on ASAN.

belitskiy committed 1 year ago

7f96b263

Skip //third_party/py/jax/tests:pytorch_interoperability_test_cpu on MSAN.

belitskiy committed 1 year ago

282ebf48

Merge pull request #22680 from superbobry:maint

a-googler committed 1 year ago

0df074c2

[Mosaic GPU] Skip matmul tests with large clusters

apaszke committed 1 year ago

d862f78d

Document `copy_to_host_async` method of jax.Array

yashk2810 committed 1 year ago

05677694

Merge pull request #22556 from cool-RR:log-cache-key

a-googler committed 1 year ago

694c14bb

Move TPU ops test to `ops_test.py`

ayaka14732 committed 1 year ago

bb160cf5

Bumped jaxlib version mypy uses on the CI

superbobry committed 1 year ago

8d33a6c9

Merge pull request #22283 from ayaka14732:ayx/lowering/sign

a-googler committed 1 year ago

2db99e03

[shard_map docs]: Fix doc typos

Google-ML-Automation committed 1 year ago

8ed94bcf

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

tlongeri committed 1 year ago

0f834cdf

Use vmap for random_gamma implementation on CPU backend

ezhulenev committed 1 year ago

15d43892

Add lowering for `lax.sign`

ayaka14732 committed 1 year ago

6cc09173

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

yashk2810 committed 1 year ago

2eb1888c

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

a-googler committed 1 year ago

6f68887e

deflake cond memory leak regression test

mattjj committed 1 year ago

f72a3f8e

Merge pull request #22674 from mattjj:remove-pdot

a-googler committed 1 year ago

f75e52db

[XLA:Mosaic] Support any int type upcast.

bythew3i committed 1 year ago

b1b7d046

Added device kwargs to jnp.linspace, jnp.array, jnp.asarray

vfdev-5 committed 1 year ago

76d61f9d

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

hawkinsp committed 1 year ago

1ac20854

remove pdot and xeinsum (since xmap is gone)

mattjj committed 1 year ago

88d1cd73

Delete mesh.Loop now that xmap has been deleted

yashk2810 committed 1 year ago

7de3c061

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

c4d3c8dd

Merge pull request #22669 from hawkinsp:repeat

a-googler committed 1 year ago

3ed9acba

Merge pull request #22668 from cool-RR:nanoseconds

a-googler committed 1 year ago

92806ee9

Added argument validation to mosaic_gpu_init_tma_desc

superbobry committed 1 year ago

5e418f5a

Newer Older