ROCm/jax

Pull Requests Commits

Consolidate AMD GPU (ROCm) installation documentation

JehandadKhan committed 16 days ago

c26f9fb6

Merge pull request #38659 from superbobry:pyrefly-2

Google-ML-Automation committed 16 days ago

1887dcbf

Drop support for Python 3.13t.

hawkinsp committed 16 days ago

60e9003f

Bumped Pyrefly to 1.1.1

superbobry committed 16 days ago

46d7e73f

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

golechwierowicz committed 16 days ago

86048571

Add --no-deps to JAX bazel wheel overrides

hawkinsp committed 16 days ago

f4c24d67

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

golechwierowicz committed 16 days ago

bdb54b97

Merge pull request #38656 from superbobry:pyrefly-2

Google-ML-Automation committed 16 days ago

abea9d1d

Merge pull request #38652 from hawkinsp:winbzl

Google-ML-Automation committed 16 days ago

a61d8cd6

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

superbobry committed 16 days ago

468355fb

Merge pull request #38653 from hawkinsp:cross

Google-ML-Automation committed 16 days ago

94a130bd

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

hawkinsp committed 16 days ago

52416ac9

Skip testSubprocessProfiling under pytest.

hawkinsp committed 16 days ago

caa31675

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

hawkinsp committed 16 days ago

416cc576

Merge pull request #38650 from hawkinsp:cross

Google-ML-Automation committed 16 days ago

1223d3be

Deprecate 2d arguments to jax.numpy.cross.

hawkinsp committed 16 days ago

d219f03b

[Mosaic GPU] Fix broken `test_kernel_with_tma`.

allanrenucci committed 16 days ago

3c146324

Merge pull request #38575 from hawkinsp:pyupgrade

Google-ML-Automation committed 17 days ago

aa4c54eb

Optimize singleton reshape sharding rule

yashk2810 committed 17 days ago

d215dea3

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

hawkinsp committed 17 days ago

faf677af

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

Google-ML-Automation committed 18 days ago

9137abb6

[trace_to_jaxpr] Refactor Pallas interpreter to use trace_to_jaxpr

gnecula committed 18 days ago

4250605d

[trace_to_jaxpr] Create `trace_to_jaxpr_nocache`

gnecula committed 18 days ago

35fe85c3

[trace_to_jaxpr] Refactor Triton lowering to use trace_to_jaxpr

gnecula committed 18 days ago

7ca0da91

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

Google-ML-Automation committed 19 days ago

2516f59d

[Mosaic GPU] Define Swizzle as Literal[32, 64, 128] | None.

golechwierowicz committed 19 days ago

92184cf1

[Mosaic GPU] Unblock the `test_get_swap_with_swizzle_transform` WG test.

golechwierowicz committed 19 days ago

0539c39d

[Mosaic GPU] Support untiled transposed 2D SMEM <-> Reg transfers.

golechwierowicz committed 19 days ago

4102d687

[Mosaic GPU] Respect user provided swizzle.

golechwierowicz committed 19 days ago

6400a754

[Mosaic GPU] Support sole swizzle in dialect lowering.

golechwierowicz committed 19 days ago

fc25aaa3

Older