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

b57f0773

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7a6f1dfe

[Mosaic] Add tpu.wait_indirect_dma assembly format

naummo committed 1 year ago

7f1934c3

Merge pull request #30132 from carlosgmartin:numpy_size_multiple_axes

Google-ML-Automation committed 1 year ago

f8e8237d

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

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

e76eb183

Avoid top-level imports in jax/_src/extend

Jake VanderPlas committed 1 year ago

da7d0ff5

Remove JAX_IFRT_VERSION_NUMBER guards after the release of 0.7.0

yashk2810 committed 1 year ago

fb1a8fcb

#sdy Delete JAX test configs enabling Shardy.

yashk2810 committed 1 year ago

ff595f61

Add LRU cache on `get_process_index_and_count`.

Hilly12 committed 1 year ago

75229a69

Avoid top-level imports in jax/_src/test_multiprocess

Jake VanderPlas committed 1 year ago

03c8a396

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

yashk2810 committed 1 year ago

430cc850

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

yashk2810 committed 1 year ago

5be0ee7b

Remove jaxlib_extension_version, ifrt_version guards now that 0.7.0 release is out

yashk2810 committed 1 year ago

5a6dd765

Reverts d86857a9c8cdcae8d27e6a371daceb0ac3f53e28

emilyfertig committed 1 year ago

8061b583

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

Google-ML-Automation committed 1 year ago

1f7373cb

Merge pull request #29541 from melissawm:refactor-httij

Google-ML-Automation committed 1 year ago

d86857a9

JAX release v0.7.0

MichaelHudgins committed 1 year ago

bfd72d5c

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

Jake VanderPlas committed 1 year ago

4b7ed577

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

Jake VanderPlas committed 1 year ago

20cf63a0

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

hyeontaek committed 1 year ago

8bfdc0dc

Refactoring HTTIJ and surrounding pages

melissawm committed 1 year ago

cce5b2a2

Extend numpy.size to accept multiple axes.

carlosgmartin committed 1 year ago

851e0716

git commit -m "Prepare for JAX release 0.7.0"

MichaelHudgins committed 1 year ago

8bc0ae6d

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

allanrenucci committed 1 year ago

0a572631

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

superbobry committed 1 year ago

9a9e2971

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

superbobry committed 1 year ago

87a4071b

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

danielsuo committed 1 year ago

eb94d52e

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

superbobry committed 1 year ago

37b6f5c5

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

MichaelHudgins committed 1 year ago

25a09917

Older