ROCm/jax

Pull Requests Commits

Bumped mypy and ruff versions used by pre-commit

superbobry committed 2 years ago

3ad1985e

Merge pull request #24437 from jakevdp:fix-diff

Google-ML-Automation committed 2 years ago

5de87808

Merge pull request #24434 from justinjfu:tutorial_mesh_update

Google-ML-Automation committed 2 years ago

1260951d

Fix jnp.diff for boolean inputs

jakevdp committed 2 years ago

66971a28

Update Pallas distributed tutorials with jax.make_mesh

justinjfu committed 2 years ago

0b46a236

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

16fca386

Rename vmap methods for callbacks.

dfm committed 2 years ago

61701af4

Merge pull request #24400 from jakevdp:subtract-ufunc

Google-ML-Automation committed 2 years ago

4a5ca2fd

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

Google-ML-Automation committed 2 years ago

65307abd

Update installation instructions

garymm committed 2 years ago

dc908b48

Make jnp.subtract a ufunc

jakevdp committed 2 years ago

6467d039

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

ecalubaquib committed 2 years ago

ad53addb

Merge pull request #24421 from jakevdp:cross-doc

Google-ML-Automation committed 2 years ago

e29b93ff

Make ffi_call return a callable

dfm committed 2 years ago

0b651f0f

Merge pull request #24417 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 2 years ago

fe83d888

Update warning message for jit of pmap

rajasekharporeddy committed 2 years ago

02f65bb1

FIx jax2tf breakge of iota

yashk2810 committed 2 years ago

783285a7

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

apaszke committed 2 years ago

f833891c

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

apaszke committed 2 years ago

f08801b8

Better docs for jnp.cross

jakevdp committed 2 years ago

a1140e92

Remove temporary aliases from .bazelrc

nitins17 committed 2 years ago

a2bc8c2e

Merge pull request #24342 from gnecula:export_custom_types

Google-ML-Automation committed 2 years ago

f4b84e1c

[export] Add support for serialization for some custom PyTree nodes

gnecula committed 2 years ago

2feea414

Merge pull request #24403 from jakevdp:load-doc

Google-ML-Automation committed 2 years ago

0d7ef9c9

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

33a73852

Remove `mesh_utils.create_device_mesh` from docs

yashk2810 committed 2 years ago

ca2d1584

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

77fb1eee

Adds arith.select to the op patters in order to canonicalize non 32 bit selects.

Google-ML-Automation committed 2 years ago

48bddc6f

Better documentation for jnp.load

jakevdp committed 2 years ago

0a85ba5f

[Pallas TPU] Use new MLIR op names

ayaka14732 committed 2 years ago

884f1dc3

Newer Older