ROCm/jax

Pull Requests Commits

[ROCm] Use clang env

Ruturaj4 committed 1 year ago

5508ed93

Update TPU test configuration tags.

hawkinsp committed 1 year ago

95f38d95

Merge pull request #23484 from justinjfu:pallas_prefetch_docs

a-googler committed 1 year ago

c0dacbf7

Merge pull request #23459 from jakevdp:split-doc

a-googler committed 1 year ago

c28b3de5

Merge pull request #23469 from rajasekharporeddy:testbranch1

a-googler committed 1 year ago

f069aebf

Slightly re-arranged Pallas Mosaic GPU pipelining logic

superbobry committed 1 year ago

05cdcb8c

Improve docs for jnp.split & related APIs

jakevdp committed 1 year ago

0320a792

Disable cudnn_fusion_test from CI.

hawkinsp committed 1 year ago

fe63b991

[pallas] Fix test failures on Windows.

hawkinsp committed 1 year ago

aa16abe5

Better dosc for jax.numpy: minimum and maximum

rajasekharporeddy committed 1 year ago

d37c8501

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

201d3ff8

Relax some test tolerances in for_loop_test.py.

hawkinsp committed 1 year ago

b6abd738

Skip symmetric product test on older jaxlibs.

hawkinsp committed 1 year ago

5af1efb2

Merge pull request #23462 from hawkinsp:mlir

a-googler committed 1 year ago

79dabe53

Adds failing test for https://github.com/google/jax/issues/23476.

jkr26 committed 1 year ago

265bb7bf

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

cd782643

Merge pull request #23490 from mattjj:readme-title

a-googler committed 1 year ago

703a8a6c

Removed dead code from `add_jaxvals`

superbobry committed 1 year ago

3e1c2b3e

tweak readme title to be more about what jax can do for you, dear user

mattjj committed 1 year ago

8ab50315

Add frontend attributes to Jax. This allows Jax users to annotate Jax code with frontend_attributes which can be traced down to the HLO level, to be used for numerical debugging purposes.

Google-ML-Automation committed 1 year ago

02b7a767

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

671acef5

[Pallas] Update Pallas docs with new figures and TPUCompilerParams

justinjfu committed 1 year ago

51a666fb

Fix tolerances for failing linalg tests.

dfm committed 1 year ago

2ce0fc25

Update FFI target name for `syrk` operation to be consistent with other kernels.

dfm committed 1 year ago

7266e338

Make pltpu key derivation more robust.

Google-ML-Automation committed 1 year ago

b6213aaa

Port GPU kernel for symmetric eigendecomposition to GPU.

dfm committed 1 year ago

1d12a993

Implement symmetric_product() to produce a symmetric matrix: `C = alpha * X @ X.T + beta * C`

Google-ML-Automation committed 1 year ago

f97bfc85

[host_callback] Fix type promotion error

gnecula committed 1 year ago

fc6b22e2

Merge pull request #23369 from sergachev:cleanup_nsys_converter

a-googler committed 1 year ago

878b6b57

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7326db77

Older