ROCm/jax

Pull Requests Commits

Add back the batching rule for ppermute

apaszke committed 5 years ago

c4f98eb8

device_put_sharded: remove incorrect type annotation

jakevdp committed 5 years ago

05cc7e73

Merge pull request #4338 from SIben:test_eig_p

a-googler committed 5 years ago

695e8d88

Ignore eig conversion test on TPU/GPU, as it is unimplemented in JAX.

bchetioui committed 5 years ago

1d94363d

Replace manual conj + transpose with call to adjoint.

bchetioui committed 5 years ago

015bc3c2

[jax2tf] Implementation of the conversion of eig_p.

bchetioui committed 5 years ago

1cde76b1

Merge pull request #4356 from google:xla-computation-dont-initialize-backend

a-googler committed 5 years ago

85d070f0

fix logic, skip test

mattjj committed 5 years ago

1092fa1d

mattjj committed 5 years ago

1aab5ced

don't initialize backend in xla_computation

mattjj committed 5 years ago

50dd9c50

Merge pull request #4343 from SIben:test_cholesky

a-googler committed 5 years ago

d9f9b50d

Merge pull request #4353 from google:xla-computation-donate-argnums

a-googler committed 5 years ago

ff5c1457

add trivial test

mattjj committed 5 years ago

a6b3fa2c

mattjj committed 5 years ago

e88579f2

Merge pull request #4330 from google:create-token-stop-grad

a-googler committed 5 years ago

cfbaca05

plumb donate_argnums into jax.xla_computation

mattjj committed 5 years ago

f172fb74

trivial change to test source sync

froystig committed 5 years ago

16e60360

Fixed ppermute translation rule (#4349)

Chase Roberts committed 5 years ago

Verified 2bc92f55

Expose scale_and_translate as a public function, fix a bug in implementation when translation is not 0.

johnpjf committed 5 years ago

7fd7009c

fix incorrect indentation

jakevdp committed 5 years ago

6b9dfb13

Make the custom_assert a one-liner.

bchetioui committed 5 years ago

b3930f0b

[jax2tf] Add primitive conversion for cholesky_p.

bchetioui committed 5 years ago

9ddd252b

rename and simplify TypedJaxpr -> ClosedJaxpr (#4328)

mattjj committed 5 years ago

Verified 6614f948

fix benchmark sums (#4329)

jakevdp committed 5 years ago

Verified 6a89f606

Enable complex-valued Cholesky decomposition tests on TPU> (#4339)

hawkinsp committed 5 years ago

Verified 2911bcd6

Disable testExpmGrad on TPU, pending investigation of compiler error (#4333)

gnecula committed 5 years ago

Verified 8376d920

[jax2tf] Revert '[jax2tf] Replace tf.math.add with tf.raw_ops.AddV2 (#4278)' (#4332)

gnecula committed 5 years ago

Verified ded7b385

[jax2tf] Replace tf.math.add with tf.raw_ops.AddV2 (#4278)

gnecula committed 5 years ago

Verified 0ac25c76

Fix conditional in eig and expand eig test suite. (#4320)

bchetioui committed 5 years ago

Verified d478e346

Use jvp(expm) to compute expm_frechet.

zhangqiaorjc committed 5 years ago

49a01d36

Newer Older