ROCm/jax

Pull Requests Commits

Update ROCm runner labels to new labels

psanal35 committed 112 days ago

20eafd5a

Merge pull request #38617 from Kinetic-Labs-GT:docs/fix-multivariate-normal-sqrt

Google-ML-Automation committed 112 days ago

339b7643

Merge pull request #38613 from RoyK0108:doc-trim-zeros-jit

Google-ML-Automation committed 112 days ago

67e967df

[Pallas:MGPU] Fix SMEM RefUnion collision across multiple `pl.run_scoped`.

allanrenucci committed 112 days ago

b7018a81

Merge pull request #38659 from superbobry:pyrefly-2

Google-ML-Automation committed 112 days ago

1887dcbf

Drop support for Python 3.13t.

hawkinsp committed 112 days ago

60e9003f

docs: fix normalization factor math formulas in multivariate_normal docstrings

Kinetic-Labs-GT committed 112 days ago

47403ab4

Bumped Pyrefly to 1.1.1

superbobry committed 112 days ago

46d7e73f

[Mosaic GPU] Remove `skip_if_wg_semantics` for `gpu_pallas_distributed_test.py`.

golechwierowicz committed 112 days ago

86048571

Document that jnp.trim_zeros is not JIT-compatible

RoyK0108 committed 112 days ago

3ad9592c

Add --no-deps to JAX bazel wheel overrides

hawkinsp committed 112 days ago

f4c24d67

[Mosaic GPU] Support duplicate 1-dim strides in `is_valid_smem_layout_assignment`.

golechwierowicz committed 112 days ago

bdb54b97

Merge pull request #38656 from superbobry:pyrefly-2

Google-ML-Automation committed 112 days ago

abea9d1d

Merge pull request #38652 from hawkinsp:winbzl

Google-ML-Automation committed 112 days ago

a61d8cd6

Use Python 3.12 for the debug-statements pre-commit hook

superbobry committed 112 days ago

468355fb

Merge pull request #38653 from hawkinsp:cross

Google-ML-Automation committed 112 days ago

94a130bd

Remove documentation for 2D jnp.cross usage, which now issues a warning.

hawkinsp committed 112 days ago

52416ac9

Skip testSubprocessProfiling under pytest.

hawkinsp committed 112 days ago

caa31675

Add workaround for https://github.com/bazel-contrib/rules_python/issues/3721

hawkinsp committed 112 days ago

416cc576

Merge pull request #38650 from hawkinsp:cross

Google-ML-Automation committed 112 days ago

1223d3be

Deprecate 2d arguments to jax.numpy.cross.

hawkinsp committed 112 days ago

d219f03b

[Mosaic GPU] Fix broken `test_kernel_with_tma`.

allanrenucci committed 112 days ago

3c146324

Merge pull request #38575 from hawkinsp:pyupgrade

Google-ML-Automation committed 112 days ago

aa4c54eb

Optimize singleton reshape sharding rule

yashk2810 committed 113 days ago

d215dea3

Split the "multiaccelerator" tag into "multiaccelerator" and "multiaccelerator-only".

hawkinsp committed 113 days ago

faf677af

Update XLA dependency to use revision http://github.com/openxla/xla/commit/3b0ff804f29ee5bb224854bf0e12fc0cbecb6db9

Google-ML-Automation committed 113 days ago

9137abb6

[trace_to_jaxpr] Refactor Pallas interpreter to use trace_to_jaxpr

gnecula committed 114 days ago

4250605d

[trace_to_jaxpr] Create `trace_to_jaxpr_nocache`

gnecula committed 114 days ago

35fe85c3

[trace_to_jaxpr] Refactor Triton lowering to use trace_to_jaxpr

gnecula committed 114 days ago

7ca0da91

Update XLA dependency to use revision http://github.com/openxla/xla/commit/6c9d76d776efd5d700845c3f36da28adbb8cd846

Google-ML-Automation committed 114 days ago

2516f59d

Older