ROCm/jax

Pull Requests Commits

[XLA][Mosaic] Add support for fp8 matmuls in TPUv5+

Google-ML-Automation committed 1 year ago

1edb94ec

Hypothesis isn't compatible with subTest and jax tests sometimes report these

pschuh committed 1 year ago

5eb679ee

Merge pull request #21568 from mattjj:custom-vjp-bwd-produce-list-or-tuple

a-googler committed 1 year ago

c8207ee8

[custom_vjp] allow bwd rule to produce top-level list (not just tuple)

mattjj committed 1 year ago

2f4cc6e9

Fix copy.deepcopy support for arrays in pinned_host memory.

yashk2810 committed 1 year ago

05916209

Merge pull request #21563 from dfm:fix-ffi-test

a-googler committed 1 year ago

4afd32ee

[XLA:Mosaic] Fix infer/apply vector layout rule for terminators (scf::yieldOp, scf::conditionOp).

bythew3i committed 1 year ago

389bf93a

Remove failing ffi test

dfm committed 1 year ago

690fa1d9

Merge pull request #21531 from dfm:move-ffi-submodule

a-googler committed 1 year ago

d9f07d03

Merge pull request #21544 from marcoselvi:patch-2

a-googler committed 1 year ago

29632fb3

Move `jax.ffi` submodule to `jax.extend.ffi`

dfm committed 1 year ago

1e206880

Pallas GPU no longer assumes that all slices have stride 1

superbobry committed 1 year ago

d6a84cc5

[jax2tf] Bump tolerance for FFT tests on CPU.

gnecula committed 1 year ago

d26819d6

[Mosaic GPU] Restore the PTX/PTXAS/SASS dump flags

apaszke committed 1 year ago

d01496a0

Test the loading of parameters from host memory

cota committed 1 year ago

ccfe01c6

[Mosaic GPU] Add a simple benchmark.

cperivol committed 1 year ago

8eaea2b1

Updated the layer norm implementation in Mosaic GPU tests

superbobry committed 1 year ago

d2a39bc6

Wrap wgmma.fence in llvm.inline_asm to constrain LLVM scheduling

apaszke committed 1 year ago

41685db0

Disable PgleTest.testPGLEProfilerGetFDOProfile.

Google-ML-Automation committed 1 year ago

33c7c8d3

Decrease tile sizes in Pallas tests

apaszke committed 1 year ago

3fb6817f

[mosaic:gpu] Minor cleanup in matmul example.

chr1sj0nes committed 1 year ago

9a572d23

[mosaic:gpu] Relax constraint on `stages` in matmul example.

chr1sj0nes committed 1 year ago

3c640660

Fix parenthesis in "Gradients contain NaN where using where"

marcoselvi committed 1 year ago

Verified 7a8bcf6e

Merge pull request #21430 from Cjkkkk:remove_is_training

a-googler committed 1 year ago

8deed95c

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

38b34b23

Make sure that the sharding and unconstrained_dims in with_sharding_constraint are correct when wsc is vmapped.

yashk2810 committed 1 year ago

54eaef9f

Make UNCONSTRAINED's `__str__` and `__repr__` so it prints nicely in all cases

yashk2810 committed 1 year ago

7aaf29bf

Improve the error message when users pass DeviceLocalLayout.AUTO to `jax.jit` and a jax.Array as an argument.

yashk2810 committed 1 year ago

bfaf0b74

Merge pull request #19538 from mattjj:vjp-mismatch-error-message

a-googler committed 1 year ago

bca9882d

Merge pull request #21071 from mattjj:vmap-spmd-axis-name-errors

a-googler committed 1 year ago

b1d37d1d

Newer Older