ROCm/jax

Pull Requests Commits

Add mxfp4 to `jax.nn.scaled_dot_general` and fix two MX quantizer defects

ScXfjiang committed 31 days ago

a27bd440

[state] Forward `jax.Ref.at` to the aval

superbobry committed 31 days ago

e4992c3d

[pallas:fuser] Skip `physicalize` when no `FusionDType` is present

superbobry committed 31 days ago

8b2a191b

[pallas:fuser] Squeeze mapped axis in `Fusible.batch` instead of using it for indexing

superbobry committed 31 days ago

545ec2c5

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

Google-ML-Automation committed 31 days ago

e9574d99

[mgpu] Add support for CTA barriers in Mosaic GPU.

chr1sj0nes committed 31 days ago

b7f91bfd

Currently the reshape layout rule always moves singleton dims to major-most. But if the singleton dim is part of the TPU tile, then moving it to major-most requires a physical copy.

yashk2810 committed 31 days ago

23d77aa0

Prepend `in_specs` for separated constants in `_shard_map_to_lojax`.

chr1sj0nes committed 31 days ago

efd0ff25

Tighten TPU float32 `expm1` (`HIGHEST` accuracy) bound from 2.0 ULP to 1.0 ULP.

hawkinsp committed 31 days ago

b02dd0ed

Improve float32 accuracy of lax.lgamma to <= 4 ULP

hawkinsp committed 31 days ago

c75467ef

Fix reshard when the dst_sharding has an empty mesh in `_reshard_shaped_array`

yashk2810 committed 31 days ago

7a6b6b18

[Pallas:TPU] Lower `lax.mulhi` in Pallas TPU and lower extended 32-bit multiplies in Mosaic.

hawkinsp committed 31 days ago

ab30a218

Add support for `jax.sharding.reshard` on `HiType`s

yashk2810 committed 31 days ago

347f313a

Merge pull request #41406 from mattjj:push-qlslnwumloxl

Google-ML-Automation committed 31 days ago

87769654

[custom_vjp] error when grad+vmap+custom_vjp may produce wrong gradients

mattjj committed 32 days ago

13f95af3

[pallas:mosaic] Deprecated `pltpu.BufferType.MANUAL` and `pltpu.BufferedRefBase`

superbobry committed 32 days ago

8fabfbc5

Set `jax_mosaic_unzip_debug_locations` to be true by default.

yueshengys committed 32 days ago

ca4c340d

[pallas:fuser] Fixed usage rules for mixed REGULAR/SCALAR_PREFETCH usage

superbobry committed 32 days ago

c0551ad9

[pallas:fuser] Support Blocked, Squeezed, and None block dims in `_swap_eval_rule`

superbobry committed 32 days ago

e0baa9fa

[XLA:CPU] Improve float32 and float64 expm1 precision from 5-6 ULP to < 1 ULP.

hawkinsp committed 32 days ago

c7a6895e

Tag `lax_numpy_test` and `lax_numpy_setops_test` as `nomsan` on CPU.

bchetioui committed 32 days ago

ba9008dc

Set `shard_count = 4` for `tpu_pallas_interpret_distributed_test` to prevent CPU timeouts.

bchetioui committed 32 days ago

fd67b6c2

Merge pull request #41508 from ROCm:feature/rocm-dpx-pytest-workers

Google-ML-Automation committed 32 days ago

8497aaa5

[pallas:fuser] Make fuser rules treat `pl.Blocked(n)` like `n`

superbobry committed 32 days ago

5390e0d1

Merge pull request #41287 from Intel-tensorflow:kranipa/oneapi_libs_runpath

Google-ML-Automation committed 32 days ago

0fa7935b

[hijax] avoid using eval_shape in Nonzero

Jake VanderPlas committed 32 days ago

3aceb9be

Merge pull request #41408 from ahlag:fix-40626-scatter-extremal-mode

Google-ML-Automation committed 32 days ago

f26f92c0

Merge pull request #41360 from Intel-tensorflow:oneapi-contact-update

Google-ML-Automation committed 32 days ago

87d9d094

Improve tail accuracy of lax.erf_inv, erfcinv, and ndtri without FMA or float64.

hawkinsp committed 32 days ago

ffdde91a

Propagate scatter mode through scatter_max/min JVP

ahlag committed 32 days ago

ae23e711

Older