ROCm/jax

Pull Requests Commits

[ROCm] Fix GPU handle pool singleton aliasing using opaque handle type separation

phambinhfin committed 1 year ago

c761d902

Update the default value of `get_concrete_mesh` from None to empty concrete mesh if there is no mesh in context. This is consistent with `get_abstract_mesh` also returning an empty abstract mesh with no mesh in context.

yashk2810 committed 1 year ago

476efd34

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

43edd643

Merge pull request #30696 from kanglant:update_py314_ci_configs

Google-ML-Automation committed 1 year ago

dbf740fb

[JAX] Move Py{Client,Array,Device,...} into jax:: namespace.

hawkinsp committed 1 year ago

036dff60

Use compression.zstd rather than zstandard under Python 3.14.

hawkinsp committed 1 year ago

732a2759

Automated Code Change

Google-ML-Automation committed 1 year ago

7ab19518

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

717b2852

Merge pull request #30697 from mattjj:boxes5

Google-ML-Automation committed 1 year ago

a89fb592

[hijax] refactor so hijax interface and boxes are in jax._src.hijax

mattjj committed 1 year ago

4fe08bf0

Merge pull request #30519 from melissawm:reorganize-tutorials

Google-ML-Automation committed 1 year ago

0abd6962

Merge pull request #30691 from gnecula:jaxpr_consts_fix_tests

Google-ML-Automation committed 1 year ago

48296a85

Reorganize the tutorials

melissawm committed 1 year ago

916658a0

Allow a PyTree future from serialization to be printable without errors

rdyro committed 1 year ago

d1abd0a2

Use `LRUCache::ForEach` in jaxlib.

mwhittaker committed 1 year ago

3bea7ff8

Fix a build failure under Python 3.14.

hawkinsp committed 1 year ago

5d08339c

Add devices fingerprint into the base launch ID. This ensures that the launch ID is unique if the same executable is executed on different subset of devices.

Google-ML-Automation committed 1 year ago

e540fe35

Accelerate deprecation of jax.experimental.pallas.tpu.TPUCompilerParams and jax.experimental.pallas.tpu.TPUMemorySpace.

danielsuo committed 1 year ago

6330af95

Correctly set local_shape to allow submesh array serialization

rdyro committed 1 year ago

bdb83318

Add python 3.14 nightly tests

kanglant committed 1 year ago

c03c92e2

Add `.update` to Mesh and AbstractMesh and do a few cleanups

yashk2810 committed 1 year ago

e2a30959

[jaxpr consts] Fix arg names when we use const args

gnecula committed 1 year ago

3c534c7e

Construct XLA GPU client with coordination service client.

mwhittaker committed 1 year ago

d5811ff7

Add Bazel TPU tests in JAX presubmit.

Google-ML-Automation committed 1 year ago

ab90f913

Temporarily disable the failing test testFrexp on Windows platform.

Google-ML-Automation committed 1 year ago

43e5b234

Move pytree code into jax:: namespace.

hawkinsp committed 1 year ago

84d3e61b

Merge pull request #30676 from mattjj:array-ref-doc-tweak

Google-ML-Automation committed 1 year ago

ef8b241c

[JAX] Rename xla::nb_class_ptr to jax::nb_class_ptr.

hawkinsp committed 1 year ago

5754a76f

[Mosaic GPU] Add a missing WG barrier before async_copy in dialect lowering + add an option to disable auto barriers

apaszke committed 1 year ago

acd84d12

Remove `jax.sharding.use_mesh` in favor of `jax.set_mesh`. `jax.set_mesh` can act as a context manager and a global setter!

yashk2810 committed 1 year ago

93f7bd0c

Older