ROCm/jax

Pull Requests Commits

Update xla reference

alekstheod committed 55 days ago

36e7372f

[Mosaic:GPU] Nicer error message for partitioned copies.

golechwierowicz committed 55 days ago

0406a74a

[Mosaic GPU] Add layout inference rule for `MemRefReshapeOp`.

allanrenucci committed 55 days ago

53b46ec4

[pallas:mgpu] Added support CP_ASYNC copies where the GMEM ref has squeezed out dimentions

superbobry committed 55 days ago

76e4ca22

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

Google-ML-Automation committed 55 days ago

30a436d6

Make @remat work nicely with jaxtyping when f is None again!

yashk2810 committed 55 days ago

e67455a1

Improve the way in which XLA optimization_level is set.

Michael Wong committed 55 days ago

48b8101f

evaluate_sampler_bias: add support for discrete distributions

Jake VanderPlas committed 55 days ago

1e944b86

Handle subclasses in AbstractRef methods

dougalm committed 55 days ago

33420d13

Remove a redundant check in dot_sharding_rule when contracting dims are sharded

yashk2810 committed 55 days ago

06818866

Reverts 43892f526e0078f823077da76b36f56583e4dd9a

Cristian Garcia committed 55 days ago

4e876607

[JAX] Adjusting mesh for allowing 2 cores per chip on V5

BrianWieder committed 55 days ago

e8a2ce08

[mosaic_gpu] Allow FP8 MMA with mixed types (E4M3FN vs E5M2)

Google-ML-Automation committed 55 days ago

657b315d

[Pallas:TPU] Add support for int8 to explicit MXU API for TPUv6e and TPUv5*.

Google-ML-Automation committed 55 days ago

bc36c9d3

Merge pull request #39400 from ROCm:reduce-rbe-jobs

Google-ML-Automation committed 55 days ago

d16e1ff6

Make logistic_p branch on tpu generation.

yashk2810 committed 56 days ago

86dca34d

Merge pull request #39401 from hawkinsp:mpmath

Google-ML-Automation committed 56 days ago

90b208c2

[ROCm] Reduce ROCm RBE test parallelism from 50 to 32 jobs

magaonka-amd committed 56 days ago

89494fbb

Merge pull request #39381 from jakevdp:lapack-dtype-error

Google-ML-Automation committed 56 days ago

7e5cb776

Remove the mpmath pin in test requirements.

hawkinsp committed 56 days ago

29b9cde9

jax.random.chisquare: support method='approximate'

Jake VanderPlas committed 56 days ago

a036ac82

[pallas] Treat `pl.debug_check` effects as kernel local

superbobry committed 56 days ago

ddf234a1

[Pallas] Adjust two erroneously updated `tpu_info_test`s

bchetioui committed 56 days ago

1e32ba12

[pallas:sc] Added a bit more validation logic to abstract eval of various primitives

superbobry committed 56 days ago

157039cc

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

Google-ML-Automation committed 56 days ago

54ba3e02

[Mosaic GPU] Parameterize interpret mode code over vector clock

Paul Biberstein committed 56 days ago

143634d4

Move tpu_info to JAX in preparation for JAX lowering rules to branch of TPU generation.

yashk2810 committed 56 days ago

9e83c445

[hijax] handle hijax primitives in roofline

Jake VanderPlas committed 56 days ago

1760967b

Error out if lhs_batch_spec and rhs_batch_spec don't match exactly.

yashk2810 committed 56 days ago

a8976b50

jax.random.beta: support method='approximate'

Jake VanderPlas committed 56 days ago

2994749a

Older