ROCm/jax

Pull Requests Commits

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

Merge pull request #23928 from nvcastet:fix_mosaic_barrier

Google-ML-Automation committed 1 year ago

a6b4648e

Improve documentation for jax.jacobian

jakevdp committed 1 year ago

cf51ee7e

[Pallas/Mosaic GPU] Replace tiling/transpose fields of GPUBlockSpec with a transform list

apaszke committed 1 year ago

3c25da2c

[pallas::mosaic_gpu] Turn the accumulator into a reference

cperivol committed 1 year ago

b6d668e0

Merge pull request #23899 from rajasekharporeddy:testbranch2

Google-ML-Automation committed 1 year ago

f6fdfb45

Better doc for jnp.trace

rajasekharporeddy committed 1 year ago

8ffeb238

Merge pull request #23924 from jburnim:jburnim_doc_typo_fix

Google-ML-Automation committed 1 year ago

911acf1b

Roll back cl/678765762 [Mosaic TPU] Support bitcast without forcing retiling.

bythew3i committed 1 year ago

e4ca4f5a

Change `CLANG_CUDA_COMPILER_PATH` set order. Add `--config=cuda_clang` to build.py

Google-ML-Automation committed 1 year ago

f1b3251b

Improve docs for jnp.pad

jakevdp committed 1 year ago

ad6c3a7f

[Mosaic GPU] Fix mbarrier inline ptx for newer CTK

nvcastet committed 1 year ago

08629a42

Fix make_remote_async_copy -> make_async_remote_copy in async doc.

jburnim committed 1 year ago

a1f2edc9

Remove `CC="/usr/lib/llvm-18/bin/clang"` from clang config in .bazelrc

Google-ML-Automation committed 1 year ago

ce99c18a

Newer Older