ROCm/jax

Pull Requests Commits

[ROCm] Align hyper parameters

amd-jianli12 committed 42 days ago

4c50a607

[ROCm] Support lowering through PJRT_Triton_Extension

amd-jianli12 committed 42 days ago

9a5fdafd

[Mosaic GPU] Add `nvshmemx_cumodule_finalize` support.

allanrenucci committed 57 days ago

2da97136

Update XLA dependency to use revision http://github.com/openxla/xla/commit/7521349eccb22a780f50fa5f6f09dbaa1d09f470

Google-ML-Automation committed 57 days ago

f29645bb

Update XLA dependency to use revision http://github.com/openxla/xla/commit/08872f587d442a05802cbdb052e8c9e6e87423f4

Google-ML-Automation committed 58 days ago

0ad46d99

Update XLA dependency to use revision http://github.com/openxla/xla/commit/66dbbf501ffd74f83c6a5d8fc201c756b1198d64

Google-ML-Automation committed 59 days ago

6469f170

Fix a broken sharded -> unreduced test

yashk2810 committed 60 days ago

20bdca7a

Add replicated -> unreduced test coverage

yashk2810 committed 60 days ago

d3ed2f4a

Support weakref_lru_cache.evict.

pschuh committed 60 days ago

dd7ff373

[dep] make jax_default_dtype_bits config a no-op for JAX v0.9.0.

Jake VanderPlas committed 60 days ago

514b7c47

Integrate LLVM at llvm/llvm-project@7d381f2a5634

Google-ML-Automation committed 60 days ago

4f0f93eb

Merge pull request #34065 from mattjj:custom-vjp3-revival

Google-ML-Automation committed 60 days ago

69518c59

[hijax] upgrade VJPHiPrimitive.def_fwd to support symbolic zeros

mattjj committed 60 days ago

ec5015fc

Merge pull request #34060 from jakevdp:interpolation-dep

Google-ML-Automation committed 60 days ago

3f8f860a

[dep] Remove deprecated jax.lib submodules for JAX v0.9.0.

Jake VanderPlas committed 60 days ago

82f02eb9

Remove pre-0.8.2 version guards.

belitskiy committed 60 days ago

1085d3a0

[dep] remove references to the already-deprecated interpolation argument

jakevdp committed 60 days ago

d6944d3d

[dep] finalize deprecation of lax.dot positional arguments for JAX v0.9.0.

Jake VanderPlas committed 60 days ago

f314722e

[dep] remove deprecated `jax_safer_randint` configuration for JAX v0.9.0

Jake VanderPlas committed 60 days ago

a9dac116

Merge pull request #34058 from jakevdp:finalize-deps

Google-ML-Automation committed 60 days ago

f29de060

Enable using custom hermetic NCCL version.

ybaturina committed 60 days ago

cfc997cd

Merge pull request #33884 from melissawm:advanced-ad

Google-ML-Automation committed 60 days ago

0d0d4913

Merge pull request #33828 from rao-ashish:asrao/cross_host_transfer_batching

Google-ML-Automation committed 60 days ago

a82d51d2

Default make_mesh to Explicit axes for the upcoming 0.9.0 release

yashk2810 committed 60 days ago

89401bd3

Remove an erroneously added license.

belitskiy committed 60 days ago

a26a5441

[dep] remove several dozen finalized deprecations for v0.9.0

jakevdp committed 60 days ago

e27c8dd1

Automated Code Change

cushon committed 60 days ago

7bc428b8

[Mosaic GPU] Add a sanity check ensuring that layout inference only ever assigns layouts compatible with variables.

bchetioui committed 60 days ago

d00f43f8

internal changes

al-bus committed 60 days ago

585016cb

[MGPU] Add WGStridedFragLayout broadcast support on major dim.

golechwierowicz committed 60 days ago

55c215d0

Older