ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 322 days ago

767c741d

[Mosaic TPU] Enable non-sublane-aligned 2D int8 load/stores

tlongeri committed 322 days ago

1d25c82a

[Pallas] Generalize BlockSpec to support different indexing mode for each dim in the block shape

sharadmv committed 323 days ago

eab1dfcc

Rename `with_user_mesh` to `with_explicit_mesh`

yashk2810 committed 323 days ago

a2ebdf6d

Enter into auto mode for `.at[...].get(...)` a bit earlier so that all ops inside `_gather` are in auto mode.

yashk2810 committed 323 days ago

7de522c5

Remove code to support jaxlib < v0.6.

hawkinsp committed 323 days ago

474dcd40

Return empty `wheel sources map` in case if `wheel_sources` is `None`.

Google-ML-Automation committed 323 days ago

62a46ada

jnp.frexp: add custom JVP rule for proper derivatives.

jakevdp committed 323 days ago

3e585c68

Merge pull request #28103 from dfm:pe-src-info

Google-ML-Automation committed 323 days ago

c2ba1790

Remove roofline subtests so they work with pytest.

zacmustin committed 323 days ago

6c4998e2

Merge pull request #28041 from ROCm:fix-clang-config-upstream

Google-ML-Automation committed 323 days ago

d45232bd

[pallas:mosaic] Replaced `device_type=` with `kernel_type` in `TPUCompilerParams`

superbobry committed 323 days ago

23c973e4

Remove unused jax_spmd_mode flag.

pschuh committed 323 days ago

7634230c

Enable cross-compilation builds of the wheels.

Google-ML-Automation committed 323 days ago

309d9d26

Don't recompute source_info for each tracer during staging.

dfm committed 323 days ago

1d652ab7

Merge pull request #27963 from justinjfu:tpu_docs_fixit

Google-ML-Automation committed 323 days ago

65d586a4

Use _make_lengths_same for explicit mode too.

yashk2810 committed 323 days ago

06ad3528

[Mosaic GPU] Align on device profiler array in smem to 8 bytes

andportnoy committed 323 days ago

25095153

Exclude running tests against the oldest supported libtpu for releases

nitins17 committed 323 days ago

9bd728e2

Make `ShardMapTracer` track `vma` instead of `rep`.

yashk2810 committed 323 days ago

85fdd1aa

Remove shard_map rewrite machinery since we have `vma` in types by default

yashk2810 committed 323 days ago

7d5c2691

BUG: return one covariance matrix per `rhs` in `polyfit`

scottstanie committed 323 days ago

21f9d236

Update XLA dependency to use revision

Google-ML-Automation committed 323 days ago

0149a32c

Merge pull request #28088 from vfdev-5:suppress-313-split_keys-dict-merge

Google-ML-Automation committed 323 days ago

42ee62ea

Automated Code Change

Google-ML-Automation committed 323 days ago

f2946ddc

Do not use `-> ...`

superbobry committed 323 days ago

4ceb4b05

Suppress race between split_keys_entry_added and dict_dict_merge in 3.13 TSAN CI

vfdev-5 committed 323 days ago

Verified 519f4909

Added `lax.axis_size` and switched all existing usage of `psum(1, ...)` to it

superbobry committed 323 days ago

c576d328

Use llvm::cast/dyn_cast/isa since alternatives are deprecated in https://github.com/llvm/llvm-project/pull/135556

Google-ML-Automation committed 323 days ago

76bb4953

Remove jax_varying_axes_in_types config and `rewrite` from `shard_map_p`

yashk2810 committed 323 days ago

82215f66

Newer Older