ROCm/jax

Pull Requests Commits

[ROCm] Use clang env

Ruturaj4 committed 1 year ago

6bc4fa15

[host_callback] Flip the JAX_HOST_CALLBACK_LEGACY flag to False

gnecula committed 2 years ago

2228115c

[pallas:mosaic_gpu] Shrink `max_concurrent_iteration` based on the total number of steps

superbobry committed 2 years ago

0cfed4ef

[host_callback] Skip test that only works in legacy mode.

gnecula committed 2 years ago

a644e23a

[Pallas/MGPU] Add support for transforms and swizzles on outputs

apaszke committed 2 years ago

98b72b17

[Pallas/MGPU] Make swizzle a Pallas transform

apaszke committed 2 years ago

7f655972

[Pallas/MGPU] Check for trivial indexers in get/swap lowering rules

apaszke committed 2 years ago

da5f2a3c

[Pallas/MGPU] Undo transforms before giving refs back to users

apaszke committed 2 years ago

cac2b8d5

[pallas:mosaic_gpu] Removed a stale TODO

superbobry committed 2 years ago

14ef2b6a

[Mosaic TPU] The previous change does not actually force the input offsets read by the rules, but simply disables all the checks. Reverting so that we at least regain the checks until we have a proper fix.

apaszke committed 2 years ago

f62941d1

Fix mutable array effects not being tracked properly

sharadmv committed 2 years ago

80f963c0

Merge pull request #23923 from carlosgmartin:ldexp_custom_jvp

Google-ML-Automation committed 2 years ago

31cb3fd3

[Pallas TPU] Add lowering for `lax.cos_p`

ayaka14732 committed 2 years ago

a24420e7

[Pallas TPU] Consolidate `OpsExtraTest` into `OpsTest`

ayaka14732 committed 2 years ago

23ce5a11

Edit implementation of jax.numpy.ldexp to get correct gradient.

carlosgmartin committed 2 years ago

65a58d62

Merge pull request #23995 from jakevdp:trapezoid-doc

Google-ML-Automation committed 2 years ago

c557db0b

[pallas:triton] Do not DCE the jaxpr in the lowering pass

superbobry committed 2 years ago

d74d3daa

[Mosaic TPU] Force offset to 0 when inferring input has offset out of the first tile.

bythew3i committed 2 years ago

4a596aee

Merge pull request #24025 from jakevdp:gradient-doc

Google-ML-Automation committed 2 years ago

cdc72787

Better docs for jnp.gradient

jakevdp committed 2 years ago

36d6bb90

Merge pull request #23505 from sergachev:fix_cudnn_fusion_test

Google-ML-Automation committed 2 years ago

3766f887

[Mosaic TPU] Canonicalize arith.select's condition to vector if other types are vector.

bythew3i committed 2 years ago

7e2f487a

Merge pull request #24024 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 2 years ago

bdae9ac7

Better doc for jnp.cbrt

rajasekharporeddy committed 2 years ago

5904fe15

Add a test for 64-bit precision of IFFT on GPU.

dfm committed 2 years ago

ff1c2ac1

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

504bc435

Simplify CI configuration.

hawkinsp committed 2 years ago

45cd77ad

Move test_aot_device_implicit_transfer to pjit_test.py

yashk2810 committed 2 years ago

203cda6f

[pallas:mosaic_gpu] Do not do `mgpu.commit_shared` if all outputs are invariant wrt sequential axes

superbobry committed 2 years ago

a046e21a

Remove forward compatibility checks from lowering of LU decomposition.

dfm committed 2 years ago

1a1e16ab

Older