ROCm/jax

Pull Requests Commits

Merge pull request #99 from ROCm/from_upstream_main

JehandadKhan committed 1 year ago

Verified 88bb4024

remove container

JehandadKhan committed 1 year ago

f40dec60

Project import generated by Copybara.

Default email committed 1 year ago

57bae648

Include the sdy MLIR dialect in jaxlib.

hawkinsp committed 1 year ago

5a1d0a6c

Replace disable_backends with enable_backends on jax_multiplatform_test.

hawkinsp committed 1 year ago

26632fd3

[Pallas/MGPU] Skip output transfers when they don't depend on sequenital dims

apaszke committed 1 year ago

5740ab3b

Run Pallas Mosaic GPU tests on internal CI

superbobry committed 1 year ago

afaf8b82

Fixed Pallas Mosaic GPU test following recent changes

superbobry committed 1 year ago

3ae48621

[Pallas:TPU] Fix lowering of convert_element_type(int32) -> bool.

Google-ML-Automation committed 1 year ago

ea86251a

Removed unused imports in `jax.experimental.mosaic.gpu.core`

superbobry committed 1 year ago

ea6ee4d7

[Pallas TPU] Add a note in the Pallas Quickstart documentation about the instructions of running the existing example on TPU

ayaka14732 committed 1 year ago

ab4590ce

Merge pull request #23853 from zhenying-liu:remat-scan

Google-ML-Automation committed 1 year ago

5a1549cc

[XLA:Mosaic][Pallas] Enable vector.ExtractOp for non-zero indices.

justinjfu committed 1 year ago

9f4e8d00

Merge pull request #23949 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 1 year ago

46dbb658

Improve docs for jax.numpy: sinh, cosh and tanh

rajasekharporeddy committed 1 year ago

7e6fa3ed

Merge pull request #23940 from jakevdp:jacobian-doc

Google-ML-Automation committed 1 year ago

6f7ad641

Merge pull request #23921 from rajasekharporeddy:testbranch4

Google-ML-Automation committed 1 year ago

96cf2b81

Deprecate jax.lib.xla_client.Device.

hawkinsp committed 1 year ago

0e082f97

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

140a8c70

[Pallas/MGPU] Skip outgoing TMA when the output is being revisited

apaszke committed 1 year ago

dd2ee8c7

[Pallas/MGPU] Implement block spec evaluation correctly

apaszke committed 1 year ago

076287fb

#sdy Add CPU targets in JAX.

bartchr808 committed 1 year ago

a3284bd8

Raise ValueError when axis1==axis2 for jnp.trace

rajasekharporeddy committed 1 year ago

6072f979

#sdy add JAX Shardy support for shard_map.

bartchr808 committed 1 year ago

e62a50cd

Add jax.errors.JaxRuntimeError as a public alias for the XlaRuntimeError class.

hawkinsp committed 1 year ago

7b53c2f3

Added support for `lax.cond_p` to Pallas Mosaic GPU lowering

superbobry committed 1 year ago

5cef547e

[Pallas/MGPU] Fix a race in the pipelining code

apaszke committed 1 year ago

0a66e2d0

[Pallas/Mosaic GPU] Implement a more comprehensive matmul kernel to see what we're still missing

apaszke committed 1 year ago

8599dbc9

Merge pull request #23927 from jakevdp:pad-doc

Google-ML-Automation committed 1 year ago

c07652fd

[Pallas/Mosaic GPU] Disable inference of sequential axis shapes

apaszke committed 1 year ago

57887732

Older