ROCm/jax

Pull Requests Commits

Add org-based guards to all caller workflows

mminutoli committed 179 days ago

df143d70

[ROCm] Add TheRock entries to CI test matrices

mminutoli committed 179 days ago

f2fac4ab

[ROCm] Standardize CI job naming, matrix schema,

mminutoli committed 179 days ago

d3e44d16

[ROCm] Version-tier S3 artifact layout

mminutoli committed 179 days ago

004845ef

Use pytrees for flattening and unflattening in MLIR lowering.

hawkinsp committed 179 days ago

b8105546

Disable MSan for pallas_mosaic_gpu_test

kanglant committed 179 days ago

c165682e

Reverts e5892751aa7eb93dc603404fa6d184f039bef14b

hawkinsp committed 179 days ago

8da70719

Merge pull request #37599 from bej9:comb

Google-ML-Automation committed 179 days ago

57b3c762

Merge pull request #37381 from Saumay:add-scipy-linalg-convolution-matrix

Google-ML-Automation committed 179 days ago

39f5a481

Merge pull request #37245 from samanklesaria:dtype_tests

Google-ML-Automation committed 179 days ago

eeada027

Add jax.scipy.linalg.convolution_matrix

Saumay committed 179 days ago

fd6e19ec

[Mosaic GPU][NFC] Refactor `Constraint` to make sure we handle all `Constant`s.

golechwierowicz committed 179 days ago

f544f99d

Make sure cases is broadcasted like `which` which takes shardings into account.

yashk2810 committed 179 days ago

e1eb2a99

Fix Sphinx duplicate target cross-reference warning for BlockSpec in Pallas docs.

Google-ML-Automation committed 179 days ago

114bf331

Add fuser block spec rules for `fold_in`, `min`, `multiple_of`

Google-ML-Automation committed 179 days ago

0f5ffa90

[Mosaic GPU][NFC] Enforce exhaustive `Constant` checks in constraints.py.

golechwierowicz committed 179 days ago

6a3f44ae

Use eval_jaxpr_p in _add_reshapes in shard_map.py

hawkinsp committed 179 days ago

da5319f5

[Mosaic GPU][NFC] Introduce `AlwaysTrue` constraint.

golechwierowicz committed 179 days ago

7b22f8df

[NFC] Rename `SMEMTiling` to `SMEMTransforms`.

allanrenucci committed 179 days ago

957dd4a1

Refactor `jax/_src/random/`

Jake VanderPlas committed 179 days ago

e9b27cd8

[Pallas Fuser] Support fusions that read and write Refs.

jburnim committed 179 days ago

3737444e

[Mosaic GPU][NFC] Split swizzle and memref transform extraction from transform attributes.

allanrenucci committed 179 days ago

f7e73d69

[pallas] Undeprecate pl.enable_debug_checks

gnecula committed 180 days ago

96301307

[Mosaic GPU] Fix assignment of `_TestCaseBase` in `gpu_pallas_distributed_test.py`.

bchetioui committed 180 days ago

556e6678

Merge pull request #37745 from hawkinsp:docs4

Google-ML-Automation committed 180 days ago

c1b05c57

Add documentation about the unsafety of calling deserialize_and_load with untrusted inputs.

hawkinsp committed 180 days ago

6ef459d2

[NFC] Remove unnecessary calls to `swizzle_and_transforms_from_transforms_attr`.

allanrenucci committed 180 days ago

e889ca50

Add `config.trace_context_names` for better cache miss messages.

gnecula committed 180 days ago

9819417d

[JAX] Enable CHLO lowering for cumulative sums on GPU again.

Google-ML-Automation committed 180 days ago

7609f610

Reverts a4be7b764876ff60d3fe2ebc11adec1d98b6ae3a

superbobry committed 180 days ago

7ca29680

Older