ROCm/jax

Pull Requests Commits

[ROCM] add missing typename keyword to work with gcc

Ruturaj4 committed 2 years ago

29a1cb76

[pallas] Add support for `unblocked` mode (without padding) in Triton lowering.

chr1sj0nes committed 2 years ago

712e638c

[Pallas] Simplify sign and erf_inv tests

ayaka14732 committed 2 years ago

93203c75

[pallas:mosaic_gpu] Basic implementation of wgmma.

cperivol committed 2 years ago

3e19a28b

Merge pull request #23837 from rajasekharporeddy:testbranch3

Google-ML-Automation committed 2 years ago

8362ab74

Merge pull request #23834 from rajasekharporeddy:testbranch2

Google-ML-Automation committed 2 years ago

63a890f2

Fix errata in block-sparse kernel tutorial.

ds-hwang committed 2 years ago

91f16419

Improve docs for jax.numpy: square, sqrt and modf

rajasekharporeddy committed 2 years ago

e976dee4

Merge pull request #23816 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 2 years ago

c05706b7

[Checkify] Add checks for shard_map.

Google-ML-Automation committed 2 years ago

6c52ddc9

Improve docs for jnp.logspace and jnp.geomspace

rajasekharporeddy committed 2 years ago

41eccd92

Added comparison operators to `mgpu.FragmentedArray`

superbobry committed 2 years ago

1256e18f

Improve docs for jax.numpy: conjugate, conj, imag and real

rajasekharporeddy committed 2 years ago

6a72c522

Added `is_signed` to `mgpu.FragmentedArray`

superbobry committed 2 years ago

f311e81c

Merge pull request #23821 from jakevdp:jnp-doc-examples

Google-ML-Automation committed 2 years ago

ba29d5a0

Updated Pallas Mosaic GPU lowering post Mosaic GPU restructuring

superbobry committed 2 years ago

653f07a7

Ignore scipy.stats._axis_nan_policy.SmallSampleWarning for LaxBackedScipyStatsTests.testMode

vam-google committed 2 years ago

21996854

[Pallas] Skip `atomic_cas` and `atomic_counter` tests on GPU in 64-bit mode

ayaka14732 committed 2 years ago

b6fe7939

[pallas:mosaic_gpu] Fragmented array debug printing.

cperivol committed 2 years ago

48c29f62

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

02994d6b

Merge pull request #23165 from 8bitmp3:jax-docs-advanced-tutorials

Google-ML-Automation committed 2 years ago

3c3bbb8a

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

ba74490e

Merge pull request #23812 from mattjj:custom-primal-tangent-dtype-helper

Google-ML-Automation committed 2 years ago

bceceaba

add jax.experimental.primal_tangent_dtype helper

mattjj committed 2 years ago

43cc70b7

Make `make_array_from_process_local_data` go via `device_put` if there is only 1 process.

yashk2810 committed 2 years ago

a2b39192

Test that jax.numpy docstrings include examples

jakevdp committed 2 years ago

aa551e66

[Pallas GPU] Enable Pallas `OpsExtraTest` in 64-bit mode

ayaka14732 committed 2 years ago

d63afd84

Add/update JAX Advanced Tutorials docs, ToC structure

8bitmp3 committed 2 years ago

0cf040c9

[Mosaic:TPU] Efficient relayout with internal scratch

bythew3i committed 2 years ago

6b93b358

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

a5336358

Older