ROCm/jax

Pull Requests Commits

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

phambinhfin committed 197 days 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 197 days ago

476efd34

Update XLA dependency to use revision

Google-ML-Automation committed 198 days ago

43edd643

Merge pull request #30696 from kanglant:update_py314_ci_configs

Google-ML-Automation committed 198 days ago

dbf740fb

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

hawkinsp committed 198 days ago

036dff60

Use compression.zstd rather than zstandard under Python 3.14.

hawkinsp committed 198 days ago

732a2759

Automated Code Change

Google-ML-Automation committed 198 days ago

7ab19518

Update XLA dependency to use revision

Google-ML-Automation committed 199 days ago

717b2852

Merge pull request #30697 from mattjj:boxes5

Google-ML-Automation committed 199 days ago

a89fb592

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

mattjj committed 199 days ago

4fe08bf0

Merge pull request #30519 from melissawm:reorganize-tutorials

Google-ML-Automation committed 199 days ago

0abd6962

Merge pull request #30691 from gnecula:jaxpr_consts_fix_tests

Google-ML-Automation committed 199 days ago

48296a85

Reorganize the tutorials

melissawm committed 199 days ago

916658a0

Allow a PyTree future from serialization to be printable without errors

rdyro committed 199 days ago

d1abd0a2

Use `LRUCache::ForEach` in jaxlib.

mwhittaker committed 199 days ago

3bea7ff8

Fix a build failure under Python 3.14.

hawkinsp committed 199 days 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 199 days ago

e540fe35

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

danielsuo committed 199 days ago

6330af95

Correctly set local_shape to allow submesh array serialization

rdyro committed 199 days ago

bdb83318

Add python 3.14 nightly tests

kanglant committed 199 days ago

c03c92e2

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

yashk2810 committed 199 days ago

e2a30959

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

gnecula committed 199 days ago

3c534c7e

Construct XLA GPU client with coordination service client.

mwhittaker committed 199 days ago

d5811ff7

Add Bazel TPU tests in JAX presubmit.

Google-ML-Automation committed 199 days ago

ab90f913

Temporarily disable the failing test testFrexp on Windows platform.

Google-ML-Automation committed 199 days ago

43e5b234

Move pytree code into jax:: namespace.

hawkinsp committed 199 days ago

84d3e61b

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

Google-ML-Automation committed 199 days ago

ef8b241c

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

hawkinsp committed 199 days ago

5754a76f

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

apaszke committed 199 days 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 199 days ago

93f7bd0c

Older