ROCm/jax

Pull Requests Commits

Ruturaj4 committed 1 year ago

cfce71de

[ROCm] improve gpu script

Ruturaj4 committed 1 year ago

644ac10c

[ROCM] Fix BUILD.bazel library source paths

Ruturaj4 committed 1 year 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 1 year ago

dd958adc

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

yueshengys committed 1 year ago

7f44edc0

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

798297af

Merge pull request #22900 from jakevdp:dep-bfloat16

a-googler committed 1 year ago

53ab5eb2

Add test for zero-sized host memory parameter

Google-ML-Automation committed 1 year ago

9074e854

Merge pull request #22649 from ROCm:ci_jax_export_harness

a-googler committed 1 year ago

aec6efb4

Merge pull request #22901 from ROCm:ci_test_harness_vmap

a-googler committed 1 year ago

cc966574

Remove enable_gpu and xla_python_enable_gpu from jax .bazelrc.

Jieying Luo committed 1 year ago

abe7982d

Skip flaky test_weight_offload_with_dp_on_output test on GPU backend.

kanglant committed 1 year ago

ae541203

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

Ruturaj4 committed 1 year ago

707cdd47

deprecate jax.lib.xla_client.bfloat16

jakevdp committed 1 year ago

a009e1cf

Merge pull request #22834 from rajasekharporeddy:testbranch1

a-googler committed 1 year ago

f67f73c3

Merge pull request #22896 from jakevdp:pin-sphinx

a-googler committed 1 year ago

799de71c

CI: pin sphinx to avoid build errors on 8.0

jakevdp committed 1 year ago

4f8c5a33

[ROCM] Fix export harness tests

Ruturaj4 committed 1 year ago

35c70fd3

Handle bool comparisons.

Google-ML-Automation committed 1 year ago

8b9ceb59

[Mosaic GPU] Profiler improvements

apaszke committed 1 year ago

209f6cd6

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

dfm committed 1 year 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 1 year ago

f255fb70

Merge pull request #22885 from jakevdp:dep-xla

a-googler committed 1 year ago

0ab4d685

Deprecate jax.lib.xla_client._xla

jakevdp committed 1 year ago

06f29bbb

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

yashk2810 committed 1 year ago

489fbc0e

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

a497fbc5

export jax.lib.xla_extension.HloModule

Jake VanderPlas committed 1 year ago

3d857b02

Reverts 0f103d33849ca017e6a199d0f79fa0d83b373995

superbobry committed 1 year ago

e416c667

Merge pull request #22608 from kaixih:fix_cuda_version_check

a-googler committed 1 year ago

c2c04e05

Fix CUDA version checks

kaixih committed 1 year ago

09b88430

Older