ROCm/jax

Pull Requests Commits

[ROCm] ci build and dockerfile changes

Ruturaj4 committed 1 year ago

d2e058ea

JAX release 0.4.35

hawkinsp committed 1 year ago

81991d87

Use libtpu releases rather than libtpu-nightly for jax[tpu].

hawkinsp committed 1 year ago

e4f3f8f0

Merge pull request #24465 from jakevdp:fix-mypy

Google-ML-Automation committed 1 year ago

1c6b0a91

Merge pull request #21524 from andportnoy:aportnoy/unknown-platform-lowering-warning

Google-ML-Automation committed 1 year ago

9a2dd19a

Merge pull request #24452 from jakevdp:insert-doc

Google-ML-Automation committed 1 year ago

1e41d5ef

Merge pull request #24467 from andportnoy:patch-2

Google-ML-Automation committed 1 year ago

1a2737b7

jakevdp committed 1 year ago

84985021

Merge pull request #24463 from jakevdp:ml-dtypes-version

Google-ML-Automation committed 1 year ago

92aa9a39

[Pallas TPU] Fix lowering for `jnp.remainder`

ayaka14732 committed 1 year ago

c60bafcc

Add back the import of `jtu` in `flash_attention.py`

andportnoy committed 1 year ago

Verified 63789849

Raise an error when registering a lowering for an unknown platform

andportnoy committed 1 year ago

2aaa108f

[Pallas TPU] Add lowerings for bf16 `jnp.ceil` and `jnp.floor` in TPU v6+

ayaka14732 committed 1 year ago

2b7b0742

Better docs for jnp.insert

jakevdp committed 1 year ago

48dd153e

Merge pull request #24412 from jakevdp:fromfunction-doc

Google-ML-Automation committed 1 year ago

2596a405

Require ml_dtypes >= 0.4.0

jakevdp committed 1 year ago

74fc7360

Better docs for jnp.fromfunction

jakevdp committed 1 year ago

7e38cbd6

Merge pull request #24442 from jakevdp:lexsort-doc

Google-ML-Automation committed 1 year ago

587832f2

Merge pull request #24425 from dfm:rename-vmap-methods

Google-ML-Automation committed 1 year ago

a2e4aff8

[pallas] Support for setting explicit backends to pallas_call.

cperivol committed 1 year ago

4f935636

[Pallas:MGPU] Add support for grid dims in GPUMesh

apaszke committed 1 year ago

2db03ba5

Reverts ebb75db8a523150c48376d15391f84380a2bb110

Google-ML-Automation committed 1 year ago

0b3f0e11

[Pallas:MGPU] Allow allocating transformed refs in run_scoped

apaszke committed 1 year ago

84a303f3

[sharding_in_types] Add `out_type` argument to `einsum` and `dot_general` to allow specifying for the output type. Right now, it only accept a `NamedSharding` but in the future we can allow a polymorphic type of: `jax.ShapeDtypeStruct | Sharding | Layout`.

yashk2810 committed 1 year ago

ebb75db8

Better documentation for jnp.lexsort

jakevdp committed 1 year ago

8800fe28

Fix documentation.

hmoraldo committed 1 year ago

5d3cac66

Respect dot algorithm spec on TPU backends.

pravnar committed 1 year ago

ad1aff09

Merge pull request #24420 from superbobry:maint-2

Google-ML-Automation committed 1 year ago

441aeebb

Move converter related tflite functions to tensorflow/lite repo

ecalubaquib committed 1 year ago

81bf6265

Merge pull request #22410 from garymm:patch-1

Google-ML-Automation committed 1 year ago

11eeff07

Older