ROCm/jax

Pull Requests Commits

[ROCm] ci build fixes

Ruturaj4 committed 1 year ago

bfd7075c

[mosaic] Use .clone() to duplicate a module, rather than printing and parsing it.

hawkinsp committed 1 year ago

bb5fbec6

[jax2tf] Disable jax2tf with non-native serialization.

gnecula committed 1 year ago

9088adda

Disable jax2tf test recently added in cl/688976685.

gnecula committed 1 year ago

0bc70bbd

Merge pull request #24505 from gnecula:jax2tf_bug

Google-ML-Automation committed 1 year ago

3823612e

[Pallas TPU] Add lowerings for scalar `absf` and `rsqrt`

ayaka14732 committed 1 year ago

5c614470

Add a jax_wheel Bazel rule to build jax pip packages

kanglant committed 1 year ago

af285959

Fix float0 behavior inside shard_map transpose under scan.

pschuh committed 1 year ago

9500bd45

Merge pull request #24511 from mattjj:improve-concreteness-error-in-remat

Google-ML-Automation committed 1 year ago

0d68a2bf

Update hermetic CUDA docs.

Google-ML-Automation committed 1 year ago

8c9dc21e

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

bd417ba6

improve concreteness error message in remat

mattjj committed 1 year ago

42311285

Remove silent data corruption runtime flags from persistent cache key.

Google-ML-Automation committed 1 year ago

afc78524

[Mosaic GPU] Use absl::StrCat instead std::string::operator+

apaszke committed 1 year ago

6634f5a3

Finish 0.4.35 release by removing dead code

yashk2810 committed 1 year ago

6c8e56f4

[jax2tf] Fixes a bad interaction between jax2tf.convert, TF, and call_tf.

gnecula committed 1 year ago

e5bbf3dc

Merge pull request #24490 from hawkinsp:searchsorted

Google-ML-Automation committed 1 year ago

644f881a

Merge pull request #24427 from kaixih:tolerence_jax_sdpa

Google-ML-Automation committed 1 year ago

c311c738

[pallas] `input_output_aliases` now only include refs which have been written to

superbobry committed 1 year ago

717467a8

[Pallas:MGPU] Treat each warpgroup as a single logical thread.

apaszke committed 1 year ago

bb2e2303

[Mosaic GPU] Query SM and PTX ISA dynamically using driver and LLVM

andportnoy committed 1 year ago

14e0f0e7

[Mosaic TPU] Use native vector tiling to load and store with untiled memref.

bythew3i committed 1 year ago

b8bacda2

Merge pull request #24487 from jakevdp:block-doc

Google-ML-Automation committed 1 year ago

df6e5e76

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

16f8958e

[Pallas TPU] Fix `OpsTest.test_elementwise` test for bf16 inputs

ayaka14732 committed 1 year ago

ea1fc65c

Perform searchsorted binary search using unsigned intermediate values.

hawkinsp committed 1 year ago

a7d71151

Improve docs for jnp.block

jakevdp committed 1 year ago

9bf1516a

Dot algorithms are now supported for all types, change the test to reflect it.

cperivol committed 1 year ago

62351585

Merge pull request #24476 from jakevdp:cov-doc

Google-ML-Automation committed 1 year ago

7556f66e

Better docs for jnp.cov & jnp.corrcoef

jakevdp committed 1 year ago

148f9d65

Older