ROCm/jax

Pull Requests Commits

Disable linalg_test::EighTinyNorm

Rahul Batra committed 2 years ago

68e321c1

[ROCm]: Add jax-triton support for ROCm

Rahul Batra committed 2 years ago

67ab07ed

[ROCm]: Disable some tests

Rahul Batra committed 2 years ago

92adc5cc

[ROCm]: Disable sparse tests

invalid-email-address committed 2 years ago

d2004bb6

[ROCm]: Update build script

Rahul Batra committed 2 years ago

52008bbe

Merge pull request #17905 from skye:version

a-googler committed 2 years ago

4cb5eeee

Update XLA commit and versions for jax 0.4.17 release

skye committed 2 years ago

b9c602b8

Merge pull request #17851 from mattjj:readme-tweak

a-googler committed 2 years ago

3d848080

Update README.md

mattjj committed 2 years ago

Verified dddbe434

[Pallas] Add support for pytrees in scalar prefetch

sharadmv committed 2 years ago

24ad445e

Fix mock_gpu_test on OSS build.

wang12tao committed 2 years ago

ee8af098

Merge pull request #17760 from superbobry:array-any

a-googler committed 2 years ago

c3e73c67

[XLA] Split --xla_detailed_logging_and_dumping debug flag into --xla_detailed_logging and --xla_enable_dumping.

hawkinsp committed 2 years ago

b2ac2deb

Update XLA dependency to use revision

a-googler committed 2 years ago

72407f6f

fix typo: device_fun(c) -> device_fun(x)

a-googler committed 2 years ago

17e259b1

[Pallas] Automatically turn mesh indices -> physical ids for remote DMAs

sharadmv committed 2 years ago

1c796c0f

Fix jax.device_put so it doesn't use tree_map for _check_sharding.

mjwillson committed 2 years ago

17d89ad1

Merge pull request #17889 from jakevdp:jnp-axis-types

a-googler committed 2 years ago

61d2016d

Add more API set up for Mock GPU client. Also clean up previous mock GPU client

wang12tao committed 2 years ago

c12929b0

Add tests to cover `PyTreeDef.flatten_up_to` error scenarios.

a-googler committed 2 years ago

d45fa224

[typing] allow Sequence inputs in several jax.numpy functions

jakevdp committed 2 years ago

2902b32e

Merge pull request #17888 from jakevdp:changelog

a-googler committed 2 years ago

c60b4aef

Add CHANGELOG entry for the jnp annotation change

jakevdp committed 2 years ago

9247a62b

[XLA: Python] Update PJRT plugin configuration mapping with bool type

sagyakwa committed 2 years ago

228e1925

Add type stubs for jax.numpy.

Jake VanderPlas committed 2 years ago

adba2f08

Merge pull request #17882 from jakevdp:fix-nightly

a-googler committed 2 years ago

0fe420e1

[Mosaic] Allow vector.shape_cast to (un)fold the sublane dim, for as long as it remains a multiple of sublane tiling

apaszke committed 2 years ago

c9851ac7

[Mosaic] Reimplement new relayout routines in C++

apaszke committed 2 years ago

77d11e4d

Disable tests for int4 support on non-TPU platforms.

hawkinsp committed 2 years ago

ae65a7cd

Update test to skip np.bitwise_count ufunc

jakevdp committed 2 years ago

568fa792

Older