ROCm/jax

Pull Requests Commits

[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

Update rules_python.patch to support Python 3.13.0 and update python 3.13 packages in JAX.

Google-ML-Automation committed 1 year ago

53668b88

Merge pull request #24172 from shuhand0:dev/shuhan/adopt_jaxlib0.4.34

Google-ML-Automation committed 1 year ago

bcb0f646

Merge pull request #24214 from dfm:api-docs-make-mesh

Google-ML-Automation committed 1 year ago

4025b27a

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

d36afe4f

Add jax.make_mesh to API docs.

dfm committed 1 year ago

1f0a04a4

Merge pull request #24167 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 1 year ago

db71965c

Merge pull request #24210 from jakevdp:nan-to-num-doc

Google-ML-Automation committed 1 year ago

c2deae8a

[Pallas TPU] Fix comparison lowering for unsigned integers

ayaka14732 committed 1 year ago

77613f21

#sdy use updated axis context when skipping `ManualComputationOp` creation.

bartchr808 committed 1 year ago

875f44c6

Better documentation for jnp.nan_to_num

jakevdp committed 1 year ago

2f798902

Newer Older