ROCm/jax

Pull Requests Commits

Support 128-lane A scale layout in Mosaic GPU layout inference.

golechwierowicz committed 3 days ago

f5b665b0

[jax-buildcop] Gate `lax.fmuladd` `stablehlo.fma` lowering on `libtpu >= 0.0.52`.

Google-ML-Automation committed 3 days ago

4d2739ca

Rename `ifrt_version` to `xla_version` in XLA and JAX.

hawkinsp committed 3 days ago

1342a37b

[Mosaic GPU] Add sm_107f test helpers and use them in Mosaic GPU tests.

golechwierowicz committed 3 days ago

ece0a47d

Add `with_logs` option to `jax.grad` and `jax.value_and_grad`.

mattjj committed 3 days ago

db828bd5

Cleanup of reduction_order + sharding-in-types

yashk2810 committed 3 days ago

2b70fbb2

[Mosaic GPU] Allow `tmem_alloc` ops of up to 576 columns on sm_107f.

golechwierowicz committed 3 days ago

55029b8f

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

Google-ML-Automation committed 3 days ago

12612cfe

sharding-in-types + reduction_order, plus cleanups

mattjj committed 3 days ago

1a1e8750

Enable jax_accurate_erf_inv by default

hawkinsp committed 3 days ago

1fc0c014

[CI] Record test source location and unparameterized name in pytest JUnit XML.

belitskiy committed 3 days ago

a32a9aa6

Merge pull request #42463 from robinbishtt:fix-kron-vander-bool

Google-ML-Automation committed 3 days ago

debd1db1

Enable more accurate default erfc in JAX.

hawkinsp committed 3 days ago

f8862307

Return signed ULP differences from `ulp_diff` and make `ulp_diff_mpmath` private.

hawkinsp committed 3 days ago

d43d446e

Fix jnp.kron and jnp.vander for boolean inputs

robinbishtt committed 3 days ago

22e530bc

Merge pull request #42467 from zazabap:dct-rfft

Google-ML-Automation committed 4 days ago

3db97979

Merge pull request #42478 from tanveer-arch:fix-choose-narrow-dtype-large-n

Google-ML-Automation committed 4 days ago

02c19267

Expose `shared_ifrt_loaded_executable()` on `PyMpmdLoadedExecutable`

Google-ML-Automation committed 4 days ago

6ddaa252

Remove deprecated keyword arguments from `jnp.clip` type signature.

Jake VanderPlas committed 4 days ago

af020627

[jax-buildcop] Eliminate redundant XLA compilations in `hyp1f1`, `hyp2f1`, and `betaln` gradient tests

Google-ML-Automation committed 4 days ago

9f8a2eaa

Merge pull request #42497 from jakevdp:bessel-doc

Google-ML-Automation committed 4 days ago

941f3e82

[jax-buildcop] Use `libstdc++` with `linux_glibc_2_35` sysroot under `--config=tsan` in `.bazelrc`

Google-ML-Automation committed 4 days ago

bbd7366b

[jax-buildcop] Enable `TF_FORCE_GPU_ALLOW_GROWTH` in `jax2tf` GPU tests

Google-ML-Automation committed 4 days ago

e352c125

[doc] improve Bessel function docs & cross-referencing

jakevdp committed 4 days ago

7bda3ab9

[jax-buildcop] Suppress spurious bfloat16 isinf warning on Windows in nan_to_num test.

Google-ML-Automation committed 4 days ago

946c7a10

Handle unknown arg_names when explaining tracing cache misses.

Google-ML-Automation committed 4 days ago

d972a82f

[jax-buildcop] Reuse compiled shapes in Bessel kn tests to avoid shard timeouts

Google-ML-Automation committed 4 days ago

37da0416

Use call_sharding_rule in SearchSorted

yashk2810 committed 4 days ago

c1e242b1

Allow mixed `Auto` meshes in `SearchSorted` and `jnp.histogram`.

levskaya committed 4 days ago

87648ea8

[Mosaic GPU][NFC] Remove the default arguments of `SmemDescriptor.slice_i64`.

allanrenucci committed 4 days ago

c08d31c1

Newer Older