ROCm/jax

Pull Requests Commits

Remove core.concrete_aval and replace with abstractify

jakevdp committed 1 year ago

3cecbf34

Fix the breakage caused by deleted enable_memories config

yashk2810 committed 1 year ago

1e221494

Delete enable_memories code in C++ since that flag is always True and cannot be turned off now.

yashk2810 committed 1 year ago

cca9afa2

fix squeeze lowering required sharding argument

shangz-ai committed 1 year ago

0b49231e

Merge pull request #25548 from andportnoy:aportnoy/mosaic-gpu-scalar-argument-64-bit-fix

Google-ML-Automation committed 1 year ago

cce40666

Don't create a mesh inside jit in sharded-computation tutorial

yashk2810 committed 1 year ago

83ca7c33

Merge pull request #25550 from hawkinsp:postrelease

Google-ML-Automation committed 1 year ago

8f4e13f6

Add Pallas Philox implementation.

justinjfu committed 1 year ago

7e96914e

[Mosaic GPU] Skip testing uint64 unless 64-bit types are enabled

andportnoy committed 1 year ago

6ea47082

Update version numbers after release.

hawkinsp committed 1 year ago

ff52aedf

Merge branch 'release/0.4.38'

hawkinsp committed 1 year ago

fd8de7d2

Merge pull request #25544 from jakevdp:array-aval

Google-ML-Automation committed 1 year ago

d4031e99

Always use the same code for array avals

jakevdp committed 1 year ago

c9afc89c

Merge pull request #25534 from jakevdp:faster-avals

Google-ML-Automation committed 1 year ago

05ad3934

Enable PJRT compatibility in cloud TPU CI.

hawkinsp committed 1 year ago

71b23ea2

Merge pull request #25538 from jakevdp:raise-runtimeerror

Google-ML-Automation committed 1 year ago

4aebe55f

Release v0.4.38

hawkinsp committed 1 year ago

c4e3f5a3

Raise rather than return error

jakevdp committed 1 year ago

2518c623

Streamline some core.concrete_aval compute paths

jakevdp committed 1 year ago

a2ac234f

Merge pull request #25456 from jakevdp:xla-abstractify

Google-ML-Automation committed 1 year ago

0fa54197

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7fe25791

Merge pull request #25473 from adamjstewart:build/no-git

Google-ML-Automation committed 1 year ago

7bbba40c

Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0

hawkinsp committed 1 year ago

7de9eb20

[Mosaic TPU] Add support for the interleaved pack format to tpu.unpack_subelements

apaszke committed 1 year ago

4911a396

[Mosaic GPU] Add end-to-end lowering example for a pointwise kernel using the dialect and layout inference.

bchetioui committed 1 year ago

36b12d58

Cleanup: toward merging core.concrete_aval & xla.abstractify

jakevdp committed 1 year ago

2c722d9b

Put abstract_mesh on every eqn so that we can preserve it during `eval_jaxpr` and `check_jaxpr` roundtrip.

yashk2810 committed 1 year ago

473e2bf5

Reduce the max_examples for splash attention tests

apaszke committed 1 year ago

e1f037bf

Make gmm TPU kernel tests significantly cheaper

apaszke committed 1 year ago

3fc23712

[Mosaic TPU] Add support for bf16 abs

apaszke committed 1 year ago

0ec902df

Newer Older