ROCm/jax

Pull Requests Commits

Add __getitem__ to backwards compatible shims.

Google-ML-Automation committed 58 days ago

c1188e74

[pallas:mosaic] Recover `memory_space` from the aval in `aval_to_ir_type`

superbobry committed 58 days ago

b9b2d4f9

[Pallas] Allow inferring the backend from the provided compiler params.

bchetioui committed 58 days ago

3349e7a1

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

Google-ML-Automation committed 58 days ago

97c557de

Automated Code Change

Google-ML-Automation committed 58 days ago

286e3583

[pallas:mosaic] Access memory spaces via `pltpu` directly

superbobry committed 58 days ago

f81c2015

[Mosaic TPU] Support reshape which folds last two dims when the last dim is not divisible by 128.

yueshengys committed 58 days ago

4751be5c

[JAX] Track the layout defaultness more precisely for arrays created using `from_dlpack`

hyeontaek committed 58 days ago

3c4dd57f

Merge pull request #34012 from mattjj:remove-primitive-batchers

Google-ML-Automation committed 58 days ago

3e389245

dont use primitive_batchers in batching.py

mattjj committed 58 days ago

b16514bd

Add platform name to xla::ifrt::Device

ICGog committed 58 days ago

fdef5952

remove batching.primitive_batchers

mattjj committed 58 days ago

01933565

Merge pull request #33752 from jakevdp:indexing-refactor

Google-ML-Automation committed 58 days ago

5e9d77f4

Update `rules_ml_toolchain` version to remove redundant `fake_nvshmem_bootstrap_uid` library from hermetic CUDA deps.

ybaturina committed 58 days ago

c593739e

[pallas] Deprecated `pltpu.ANY` in favor of `pl.ANY`

superbobry committed 58 days ago

7b63d9d3

Fix ct_check to account for `None` cotangents too

yashk2810 committed 58 days ago

254918c8

[CI] Modify TPU v7x jobs to include bazel, and exclude some python versions

MichaelHudgins committed 58 days ago

579009c4

Add --dist=loadfile to pytest command in run_pytest_tpu.sh.

belitskiy committed 58 days ago

44085743

[Pallas MGPU] Adding support for squeezed block dims in the pipeline BlockSpecs. They can be identified with a `None` or `pl.Squeezed`.

Rifur13 committed 58 days ago

4c671ca7

Add `concrete_mesh` to reshard_p in pallas lowering rules which was missing

yashk2810 committed 59 days ago

44f67b03

Skip async store reduction tests for int64/uint64 when x64 is disabled.

belitskiy committed 59 days ago

d3a80521

Causing segfaults in python_callback_test

belitskiy committed 59 days ago

86116f8f

Export pyproject.toml in the BUILD file

superbobry committed 59 days ago

39e6ee1e

Propagate effects in the abstract eval rule for custom_vmap_p.

Google-ML-Automation committed 59 days ago

20406125

Skip test_scalar_debug_check in tpu_sparsecore_pallas_debug_check_test.

belitskiy committed 59 days ago

7f7b35e5

Reverts 58fbb8a3b7d4ecb860c4cf678cae1dd6f1079846

gnecula committed 59 days ago

b44a511b

[Pallas/interpreter] Refactor re-useable functionality out of the main source file for the TPU kernel interpreter.

Google-ML-Automation committed 59 days ago

b8d37576

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

Google-ML-Automation committed 59 days ago

e63d2a49

Use `to_cotangent_aval()` in `SymbolicZero` check in _flatten_bwd

yashk2810 committed 59 days ago

9630a2b0

Dispatch to `reshard` in `broadcast_in_dim` if only sharding is changing and everything else is the same as operand.

yashk2810 committed 59 days ago

6940903c

Newer Older