ROCm/jax

Pull Requests Commits

Fix rel_entr behavior at boundary value

sh0416 committed 1 year ago

818e7d92

[Mosaic GPU] Disable XLA autotuning in `matmul_test.py`.

bchetioui committed 1 year ago

72b111af

[Mosaic] Expand vector.shape_cast no-op detection for expanding/shrinking lane shape casts

tlongeri committed 1 year ago

97f9a5e8

[Mosaic] Expand support of vector.broadcast

tlongeri committed 1 year ago

3fb9acf0

Merge pull request #21448 from gnecula:exp_simpl1

a-googler committed 1 year ago

e8a11130

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

149a771b

[export] Simplify construction of shardings for the VJP

gnecula committed 1 year ago

7dbab168

Add `jax.make_array_from_process_local_data` to the docs. Preview: https://jax--21444.org.readthedocs.build/en/21444/_autosummary/jax.make_array_from_process_local_data.html#jax.make_array_from_process_local_data

yashk2810 committed 1 year ago

08b1cef0

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

24509f32

jnp.take: fix annotation for fill_value

jakevdp committed 1 year ago

0ff0d7b9

Merge pull request #21433 from mattjj:scan-avoid-singleton-dim-in-lowering

a-googler committed 1 year ago

5d5ce1c9

avoid singleton dim in scan lowering when unroll==1

mattjj committed 1 year ago

a24b7380

Merge pull request #21414 from mattjj:pjit-forwarding-rule

a-googler committed 1 year ago

787b7c26

add pjit forwarding rule

mattjj committed 1 year ago

0a693faf

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

9595f140

Minor comment correction.

datenglin committed 1 year ago

c29cff70

Enable a lot of previously disabled tests now that compute offload works on TPU.

yashk2810 committed 1 year ago

c577087b

Add JAX version to TPU_ML_PLATFORM_VERSION environment variable.

Google-ML-Automation committed 1 year ago

93170d9c

[Pallas][Mosaic] Add lowering rules for PRNG ops.

justinjfu committed 1 year ago

683ca2cd

[xla:python] Add support for registering custom call targets for all XLA execution stages and for XLA FFI traits

ezhulenev committed 1 year ago

d5c7ccc7

Update cuDNN to version 9.1.1 in JAX

beckerhe committed 1 year ago

15a19854

Merge pull request #21262 from vfdev-5:depr-change-ddof-to-correction-21088

a-googler committed 1 year ago

bab7f40d

Added correction arg in jnp.var and jnp.std

vfdev-5 committed 1 year ago

55f8284e

Add the `threefry_partitionable` config to `JaxprEqnContext` to allow setting it inside jit.

yashk2810 committed 1 year ago

972fd665

Merge pull request #21362 from vfdev-5:fix-array-api-compat-var-nan-21330

a-googler committed 1 year ago

9b0d4030

Merge pull request #21423 from hawkinsp:cumsum

a-googler committed 1 year ago

31d4a1ea

Merge pull request #21406 from superbobry:maint-too

a-googler committed 1 year ago

d910c5cb

Another attempt to land #20445

superbobry committed 1 year ago

15b974c9

Bump up memory requirements for pmap_test on TPU.

hawkinsp committed 1 year ago

1e01fa7b

Add note to release notes about #21403.

hawkinsp committed 1 year ago

441ab58a

Newer Older