ROCm/jax

Pull Requests Commits

Fix a bug where doing a CPU jax.Array -> TPU jax.Array transfer was not setting the layout correctly on the output TPU array.

yashk2810 committed 156 days ago

b57f0773

Update XLA dependency to use revision

Google-ML-Automation committed 156 days ago

7a6f1dfe

[Mosaic] Add tpu.wait_indirect_dma assembly format

naummo committed 156 days ago

7f1934c3

Merge pull request #30132 from carlosgmartin:numpy_size_multiple_axes

Google-ML-Automation committed 156 days ago

f8e8237d

Avoid top-level imports in jax/_src/internal_test_util/*

Jake VanderPlas committed 156 days ago

a350bde9

Add `portpicker` to the list of common requirements and update `requirements_lock_3_13_ft.txt` and `requirements_lock_3_14_ft.txt` files.

Google-ML-Automation committed 156 days ago

e76eb183

Avoid top-level imports in jax/_src/extend

Jake VanderPlas committed 156 days ago

da7d0ff5

Remove JAX_IFRT_VERSION_NUMBER guards after the release of 0.7.0

yashk2810 committed 156 days ago

fb1a8fcb

#sdy Delete JAX test configs enabling Shardy.

yashk2810 committed 156 days ago

ff595f61

Add LRU cache on `get_process_index_and_count`.

Hilly12 committed 156 days ago

75229a69

Avoid top-level imports in jax/_src/test_multiprocess

Jake VanderPlas committed 156 days ago

03c8a396

Remove `vmap` and `pmap` mentions in the rst section title.

yashk2810 committed 156 days ago

430cc850

Add `smap` to the API docs on the website. Fixes https://github.com/jax-ml/jax/issues/30407

yashk2810 committed 156 days ago

5be0ee7b

Remove jaxlib_extension_version, ifrt_version guards now that 0.7.0 release is out

yashk2810 committed 156 days ago

5a6dd765

Reverts d86857a9c8cdcae8d27e6a371daceb0ac3f53e28

emilyfertig committed 156 days ago

8061b583

Merge pull request #30404 from jax-ml:postrelease-0.7.0

Google-ML-Automation committed 156 days ago

1f7373cb

Merge pull request #29541 from melissawm:refactor-httij

Google-ML-Automation committed 156 days ago

d86857a9

JAX release v0.7.0

MichaelHudgins committed 156 days ago

bfd72d5c

Avoid top-level imports in jax/_src/test_util.py

Jake VanderPlas committed 156 days ago

4b7ed577

[jax/BUILD] avoid `_src` dependencies on top-level `jax` packages

Jake VanderPlas committed 156 days ago

20cf63a0

[JAX] Make `util.lru_cache(..., trace_context_in_key=False)` always cache

hyeontaek committed 156 days ago

8bfdc0dc

Refactoring HTTIJ and surrounding pages

melissawm committed 156 days ago

cce5b2a2

Extend numpy.size to accept multiple axes.

carlosgmartin committed 156 days ago

851e0716

git commit -m "Prepare for JAX release 0.7.0"

MichaelHudgins committed 156 days ago

8bc0ae6d

Enforce the shape of MMA operands and result in their types.

allanrenucci committed 156 days ago

0a572631

[pallas:mosaic_gpu] Do not use `lax.zeros_like_array`

superbobry committed 156 days ago

9a9e2971

[pallas] `pl.loop` now accepts `step=`

superbobry committed 156 days ago

87a4071b

Update deprecation schedule for `mlir.custom_call`, attributes of `xla_client`, and `xla_bridge.get_backend`.

danielsuo committed 156 days ago

eb94d52e

[pallas:mosaic_gpu] `async_load_tmem` no longer accepts an `idx`

superbobry committed 156 days ago

37b6f5c5

Rollback XLA archival change, current automation does not calculate the shasum correctly

MichaelHudgins committed 156 days ago

25a09917

Older