ROCm/jax

Pull Requests Commits

[JAX] bazel build rocm changes

Ruturaj4 committed 1 year ago

2b408701

[ROCm] Use clang env

Ruturaj4 committed 1 year ago

09455e40

[host_callback] Flip the JAX_HOST_CALLBACK_LEGACY flag to False

gnecula committed 1 year ago

2228115c

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

superbobry committed 1 year ago

0cfed4ef

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

gnecula committed 1 year ago

a644e23a

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

apaszke committed 1 year ago

98b72b17

[Pallas/MGPU] Make swizzle a Pallas transform

apaszke committed 1 year ago

7f655972

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

apaszke committed 1 year ago

da5f2a3c

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

apaszke committed 1 year ago

cac2b8d5

[pallas:mosaic_gpu] Removed a stale TODO

superbobry committed 1 year 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 1 year ago

f62941d1

Fix mutable array effects not being tracked properly

sharadmv committed 1 year ago

80f963c0

Merge pull request #23923 from carlosgmartin:ldexp_custom_jvp

Google-ML-Automation committed 1 year ago

31cb3fd3

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

ayaka14732 committed 1 year ago

a24420e7

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

ayaka14732 committed 1 year ago

23ce5a11

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

carlosgmartin committed 1 year ago

65a58d62

Merge pull request #23995 from jakevdp:trapezoid-doc

Google-ML-Automation committed 1 year ago

c557db0b

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

superbobry committed 1 year ago

d74d3daa

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

bythew3i committed 1 year ago

4a596aee

Merge pull request #24025 from jakevdp:gradient-doc

Google-ML-Automation committed 1 year ago

cdc72787

Better docs for jnp.gradient

jakevdp committed 1 year ago

36d6bb90

Merge pull request #23505 from sergachev:fix_cudnn_fusion_test

Google-ML-Automation committed 1 year ago

3766f887

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

bythew3i committed 1 year ago

7e2f487a

Merge pull request #24024 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 1 year ago

bdae9ac7

Better doc for jnp.cbrt

rajasekharporeddy committed 1 year ago

5904fe15

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

dfm committed 1 year ago

ff1c2ac1

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

504bc435

Simplify CI configuration.

hawkinsp committed 1 year ago

45cd77ad

Move test_aot_device_implicit_transfer to pjit_test.py

yashk2810 committed 1 year ago

203cda6f

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

superbobry committed 1 year ago

a046e21a

Older