ROCm/jax

Pull Requests Commits

Update documentation and add setup.py pypi bindings

Ruturaj4 committed 1 year ago

fefebebe

Merge pull request #25595 from jakevdp:mv-shaped-abstractify

Google-ML-Automation committed 1 year ago

76805325

[Mosaic TPU] Guard tests for new features by the libtpu version

apaszke committed 1 year ago

ad00ec1d

[Pallas:TPU] Use self.pallas_call to properly handle interpret mode

apaszke committed 1 year ago

45159494

[Mosaic GPU] Add a new tiled layout, optimized for upcasting before WGMMA

apaszke committed 1 year ago

006c65d8

[Mosaic GPU] Replace the dialect's layout enum with layouts holding the proper

bchetioui committed 1 year ago

66ad2082

[Mosaic:TPU][NFC] Small cleanup of extui rule in apply-vector-layout

tlongeri committed 1 year ago

8188c574

[Mosaic GPU][NFC] Move up the definition of `ThreadSemantics`.

dimitar-asenov committed 1 year ago

16f0f97f

Add lax.composite primitive

ghpvnist committed 1 year ago

d206cc3b

Refactor: move shaped_abstractify to core

jakevdp committed 1 year ago

676070f4

[Pallas] Fix lowering tests for reduction ops

ayaka14732 committed 1 year ago

613a0cde

Merge pull request #25593 from mattjj:ref-errors-4

Google-ML-Automation committed 1 year ago

f65ecedd

Account for tokens in `allow_spmd_sharding_propagation_to_parameters` and `allow_spmd_sharding_propagation_to_output` compile options

yashk2810 committed 1 year ago

9041b02d

Merge pull request #25596 from mattjj:whitespace

Google-ML-Automation committed 1 year ago

46b18d27

add mutable array ref error checks to scan

mattjj committed 1 year ago

e5285626

[Mosaic TPU] Support i32 vector multi reduction except cross lane.

bythew3i committed 1 year ago

3a5c4da4

[Mosaic] Improve error verbosity of tpu.memref_slice verification

naummo committed 1 year ago

6bcec910

remove whitespace from distributed arrays doc

mattjj committed 1 year ago

c6576827

Re-land changes from https://github.com/jax-ml/jax/pull/25555

Jake VanderPlas committed 1 year ago

89a54a9e

[Mosaic] Verify that the target IDs are provided in remote DMAs

naummo committed 1 year ago

de359f5c

[sharding_in_types] Check out_avals with mesh context too. This is because users can pass their own shardings to functions like `einsum`, `reshape, `broadcast`, etc

yashk2810 committed 1 year ago

af63e443

[Mosaic:TPU][NFC] Delete unused functions

tlongeri committed 1 year ago

13e721a2

[Pallas] Add version guard for non-32-bit selection in test and fix github build failure.

bythew3i committed 1 year ago

74eca134

[Mosaic GPU][NFC] Split LaunchContext into a separate file

dimitar-asenov committed 1 year ago

1ed5fa6b

[Mosaic TPU] Support direct cast i8 vector to mask

bythew3i committed 1 year ago

bf692efb

fix bug with jax.remat static_argnums not supporting int

mattjj committed 1 year ago

bb16d5ac

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

53bff861

Remove jax_enable_memories config flag. It defaulted to True for a very long time and it's time to remove the flag.

yashk2810 committed 1 year ago

8b734808

Merge pull request #25490 from jax-ml:more-linearize-fixes

Google-ML-Automation committed 1 year ago

8abb1a73

[JAX] Fix a small bug if shardings is tuple.

maxwillzq committed 1 year ago

72e5ca95

Older