ROCm/jax

Pull Requests Commits

support multi-host pmap with omnistaging

jekbradbury committed 5 years ago

Verified 491ed1e0

Canonicalize result dtype to fix double precision problem in ldexp (#4069)

dav0dea committed 5 years ago

Verified 1316562b

[jax2tf] avoid import errors when omnistaging is enabled (#4072)

gnecula committed 5 years ago

Verified 1dbdaac7

support axis_index_groups in psum(const) (#4070)

jekbradbury committed 5 years ago

Verified 9ab07d85

Add `in_parts` and `out_parts` optional arguments `jax.xla_computation`. (#4055)

Ryan Sepassi committed 5 years ago

Verified 394a33c8

fixes for pytype (#4068)

jakevdp committed 5 years ago

Verified 8923bab5

Note flush-denormal-to-zero behavior in nextafter. (#4049)

mathisgerdes committed 5 years ago

Verified b44a1bac

Make it possible to override raise_to_shaped for new types (#4064)

jakevdp committed 5 years ago

Verified c311fb77

Initial version of vmap collectives (#4005)

apaszke committed 5 years ago

Verified b75bae64

Use pytree from xla_client. (#4063)

jblespiau committed 5 years ago

Verified d8de6b61

[jax2tf] Fix bfloat16 failures on CPU/GPU with latest tf-nightly. (#4060)

bchetioui committed 5 years ago

Verified bd14f234

[jax2tf] Enable bfloat16 tests (#4057)

bchetioui committed 5 years ago

Verified aea64e8b

[jax2tf] Added note about nextafter support in TensorFlow. (#4056)

bchetioui committed 5 years ago

Verified 41f1eb41

[jax2tf] Refactor test_binary_elementwise. (#4043)

bchetioui committed 5 years ago

Verified 170b63a1

add placeholder disable_omnistaging method (#4054)

mattjj committed 5 years ago

Verified 822344d6

add scan unrolling to a previous changelog entry

froystig committed 5 years ago

df6a3da4

Update auto jaxlib install to use nvcc CUDA version instead of nvidia-smi. (#4050)

skye committed 5 years ago

Verified 428217b1

[jax2tf] Clean up test_qr and disable (temporarily) some tests for shirt_right (#4044)

gnecula committed 5 years ago

Verified 7cf236f8

[jax2tf] Change np.array dtype to jnp.float32 when target dtype is bfloat16. (#3959)

bchetioui committed 5 years ago

Verified 3458f1fc

[host_callback] Fix handling of params[donated_invars] (#4040)

gnecula committed 5 years ago

Verified 206d2f98

remove auto-parallelization transformation

froystig committed 5 years ago

34f90f55

Fix broadcasting in random.uniform and randint. (#4035)

ethanluoyc committed 5 years ago

Verified 6e4ec7cb

Enable s8/u8/s16/u16 types on TPU in tests. (#4032)

hawkinsp committed 5 years ago

Verified a169743f

[host_callback] Add support for custom gradients inside control-flow (#4029)

gnecula committed 5 years ago

Verified ece05d28

skip more ode tests on gpu, b/c slow to compile (#4028)

mattjj committed 5 years ago

Verified c564aca7

Fix type promotion in np.clip (#4023)

jakevdp committed 5 years ago

Verified 5fcca184

cumulative jet rules (#4000)

thisiscam committed 5 years ago

Verified fe9f264b

Fix typos in api.py docstrings (#4021)

j-towns committed 5 years ago

Verified 1e077129

Implement the invertible decorator in terms of custom_vjp (#3957)

apaszke committed 5 years ago

Verified cd54dd97

[jax2tf] Disable the CI tests for jax2tf. (#4019)

gnecula committed 5 years ago

Verified 0b99ca89

Older