ROCm/jax

Pull Requests Commits

[ROCm] improve gpu script

Ruturaj4 committed 2 years ago

644ac10c

[ROCM] Fix BUILD.bazel library source paths

Ruturaj4 committed 2 years ago

a2d79936

Add `mesh_shape` to the lowering context. This is to allow custom partitioning to not depend on the mesh context manager to return NamedShardings even if the arguments have NamedShardings on them.

yashk2810 committed 2 years ago

dd958adc

Change log level of clearing JAX backend caches from info to debug.

yueshengys committed 2 years ago

7f44edc0

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

798297af

Merge pull request #22900 from jakevdp:dep-bfloat16

a-googler committed 2 years ago

53ab5eb2

Add test for zero-sized host memory parameter

Google-ML-Automation committed 2 years ago

9074e854

Merge pull request #22649 from ROCm:ci_jax_export_harness

a-googler committed 2 years ago

aec6efb4

Merge pull request #22901 from ROCm:ci_test_harness_vmap

a-googler committed 2 years ago

cc966574

Remove enable_gpu and xla_python_enable_gpu from jax .bazelrc.

Jieying Luo committed 2 years ago

abe7982d

Skip flaky test_weight_offload_with_dp_on_output test on GPU backend.

kanglant committed 2 years ago

ae541203

[ROCM] Fix hipsolverSsyevd tests due to align with the rocm behavior.

Ruturaj4 committed 2 years ago

707cdd47

deprecate jax.lib.xla_client.bfloat16

jakevdp committed 2 years ago

a009e1cf

Merge pull request #22834 from rajasekharporeddy:testbranch1

a-googler committed 2 years ago

f67f73c3

Merge pull request #22896 from jakevdp:pin-sphinx

a-googler committed 2 years ago

799de71c

CI: pin sphinx to avoid build errors on 8.0

jakevdp committed 2 years ago

4f8c5a33

[ROCM] Fix export harness tests

Ruturaj4 committed 2 years ago

35c70fd3

Handle bool comparisons.

Google-ML-Automation committed 2 years ago

8b9ceb59

[Mosaic GPU] Profiler improvements

apaszke committed 2 years ago

209f6cd6

Re-land FFI port of GPU LU decomposition after fixing XLA FFI memory leak.

dfm committed 2 years ago

23da11b6

Async dispatch expensive computations on the JAX CPU backend. By setting `jax.config.update('jax_cpu_enable_async_dispatch', False)`, one could opt out of the change and recover the old behavior.

yueshengys committed 2 years ago

f255fb70

Merge pull request #22885 from jakevdp:dep-xla

a-googler committed 2 years ago

0ab4d685

Deprecate jax.lib.xla_client._xla

jakevdp committed 2 years ago

06f29bbb

Add a test for streaming in closed over constants from host to device

yashk2810 committed 2 years ago

489fbc0e

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

a497fbc5

export jax.lib.xla_extension.HloModule

Jake VanderPlas committed 2 years ago

3d857b02

Reverts 0f103d33849ca017e6a199d0f79fa0d83b373995

superbobry committed 2 years ago

e416c667

Merge pull request #22608 from kaixih:fix_cuda_version_check

a-googler committed 2 years ago

c2c04e05

Fix CUDA version checks

kaixih committed 2 years ago

09b88430

Merge pull request #22836 from superbobry:maint-2

a-googler committed 2 years ago

0f103d33

Older