ROCm/jax

Pull Requests Commits

[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

Bumped mypy and ruff versions used by pre-commit

superbobry committed 1 year ago

3ad1985e

Merge pull request #24437 from jakevdp:fix-diff

Google-ML-Automation committed 1 year ago

5de87808

Merge pull request #24434 from justinjfu:tutorial_mesh_update

Google-ML-Automation committed 1 year ago

1260951d

Fix jnp.diff for boolean inputs

jakevdp committed 1 year ago

66971a28

Update Pallas distributed tutorials with jax.make_mesh

justinjfu committed 1 year ago

0b46a236

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

16fca386

Rename vmap methods for callbacks.

dfm committed 1 year ago

61701af4

Merge pull request #24400 from jakevdp:subtract-ufunc

Google-ML-Automation committed 1 year ago

4a5ca2fd

Merge pull request #24370 from dfm:ffi-call-to-callable

Google-ML-Automation committed 1 year ago

65307abd

Update installation instructions

garymm committed 1 year ago

dc908b48

Make jnp.subtract a ufunc

jakevdp committed 1 year ago

6467d039

Move out mnist py/Jax tensorflow lite tests to tensorflow lite repo

ecalubaquib committed 1 year ago

ad53addb

Merge pull request #24421 from jakevdp:cross-doc

Google-ML-Automation committed 1 year ago

e29b93ff

Make ffi_call return a callable

dfm committed 1 year ago

0b651f0f

Merge pull request #24417 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 1 year ago

fe83d888

Update warning message for jit of pmap

rajasekharporeddy committed 1 year ago

02f65bb1

FIx jax2tf breakge of iota

yashk2810 committed 1 year ago

783285a7

[Pallas:MGPU] Add support for passing in WGMMA lhs from registers

apaszke committed 1 year ago

f833891c

[Pallas:MGPU] Allow indexing to appear anywhere in the list of transforms

apaszke committed 1 year ago

f08801b8

Newer Older