ROCm/jax

Pull Requests Commits

[ROCm] Bring up clang support for JAX+XLA

Ruturaj4 committed 1 year ago

33bcd0cb

Update Eigh kernel on GPU to use 64-bit interface when it is available.

dfm committed 1 year ago

6625a2b3

Merge pull request #24193 from dfm:vectorized-changelog

Google-ML-Automation committed 1 year ago

9d44d723

Fix listing of `vectorized` deprecation in changelog.

dfm committed 1 year ago

f55141ef

Refactor JAX's .bazelrc

nitins17 committed 1 year ago

5132a188

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

f8022e78

[sharding_in_types] Normalize partition specs when creating avals so that P(None, None) and P() are treated as replicated and equivalent. Shardings on avals are always normalized.

yashk2810 committed 1 year ago

8ef41a6e

Reenable some test cases that were disabled due to bugs that now seem fixed.

hawkinsp committed 1 year ago

66f52689

Move additional CI enabled/disabled configurations into jax BUILD files.

hawkinsp committed 1 year ago

19dbff53

Deprecate jax.lib.xla_client.PaddingType.

hawkinsp committed 1 year ago

aa3254d7

[pallas:mosaic_gpu] `pl.run_scoped` now supports scoped barriers

superbobry committed 1 year ago

70ee8e11

Add lax.FftType.

hawkinsp committed 1 year ago

94abaf43

Remove dead ducc_fft code.

hawkinsp committed 1 year ago

cf5f1577

Merge pull request #24185 from superbobry:docs

Google-ML-Automation committed 1 year ago

ddf85244

Update CI-Build job to not reference core count as part of job name

MichaelHudgins committed 1 year ago

2b4a3af0

Merge pull request #24158 from superbobry:maint-2

Google-ML-Automation committed 1 year ago

82156a1a

Use the current minimum jaxlib version for type checking on the CI

superbobry committed 1 year ago

ec745f48

[Mosaic] Parameterize the number of lanes and sublanes in TPU dialects.

Google-ML-Automation committed 1 year ago

81a95f78

[pallas] Added API docs for Triton and Mosaic GPU backends

superbobry committed 1 year ago

46e65b59

[sharding_in_types] Add support for nary ops to propagate sharding when 1 input is sharded and all others are replicated.

yashk2810 committed 1 year ago

351187d9

Add Python 3.13.0 to JAX Docker images with CUDA 12.3 and CUDA 12.1.

Google-ML-Automation committed 1 year ago

b65be4e1

[Pallas] Add lowering for threefry PRNG.

justinjfu committed 1 year ago

73418427

[Pallas TPU] Add lowerings for `lax.population_count_p` and `lax.clz_p`

ayaka14732 committed 1 year ago

3fc4ba29

[Mosaic TPU] Change getLayout to force offset to 0 when inferring input has offset out of the first tile.

bythew3i committed 1 year ago

f52b016d

Merge pull request #24218 from andportnoy:aportnoy/use-py_deps

Google-ML-Automation committed 1 year ago

64a75745

Use py_deps("absl/testing") instead of //third_party/py/absl/testing

andportnoy committed 1 year ago

2c731320

Merge pull request #24130 from MichaelHudgins:workflow-new-runners

Google-ML-Automation committed 1 year ago

62459db0

[Pallas TPU] Raise `NotImplementedError` when casting to 64-bit integer

ayaka14732 committed 1 year ago

df6a0fbd

Update ci-build to use new runners

MichaelHudgins committed 1 year ago

a7b3a05f

[Mosaic TPU][NFC] Refactor tpu matmul rule.

bythew3i committed 1 year ago

f96c5661

Older