ROCm/jax

Pull Requests Commits

[ROCm] Added a back compatibility test for GEQRF on ROCm devices.

tsrw2048 committed 17 days ago

9296bbc0

Add `scratch_shapes` to Pallas `pl.core_map`.

IvyZX committed 17 days ago

5f12495b

Add support for TPU7 device kind to mesh_utils

Google-ML-Automation committed 18 days ago

c1a26b59

Export `lower_with_sharding_in_types` via jax.extend.mlir.

Google-ML-Automation committed 18 days ago

9d220373

Removed a few unused functions

superbobry committed 18 days ago

05509540

Merge pull request #35000 from ezhulenev:no-ffi-api-1

Google-ML-Automation committed 18 days ago

38657b5d

Merge pull request #35002 from jakevdp:fix-array-methods

Google-ML-Automation committed 18 days ago

f18105fc

[bug] fix undefined name in array.argpartition

jakevdp committed 18 days ago

7519619d

[jax] Do not use ffi_api.h from XLA

ezhulenev committed 18 days ago

455ff576

[pallas mosaic] Add support for TPU7 device kind.

Google-ML-Automation committed 18 days ago

de06517d

Prepare JAX test wrappers for rules_python 1.8.3 upgrade

ecalubaquib committed 18 days ago

ebd5249d

Merge pull request #34973 from jakevdp:more-pyrefly-fixes

Google-ML-Automation committed 18 days ago

c57f2524

Upgrade to CCCL 3.2.0 with better compatibility to CUDA 12.9, 13.0, and 13.1+.

ybaturina committed 18 days ago

3694b2be

[pallas] Pulled `default_memory_space` into `Mesh`

superbobry committed 18 days ago

d55be985

[typing] fix a number of pyrefly typechecking errors

jakevdp committed 18 days ago

002ac9bd

Merge pull request #34991 from superbobry:pyrefly

Google-ML-Automation committed 18 days ago

9f1399e5

Remove deprecated coordination service option.

mwhittaker committed 18 days ago

747397e1

Bump up shard counts for api_test, and scaled_dot_test, mostly for TSAN.

belitskiy committed 18 days ago

068ae53b

[pallas] Deduplicated abstract_eval and typecheck rules for `core_map_p`

superbobry committed 18 days ago

316ffe98

[Mosaic GPU][NFC] Minor code simplification in `mosaic_gpu_ext.cc`.

bchetioui committed 18 days ago

41c6e2a8

[pallas:mosaic] Extracted `lower_jaxpr_into_module`

superbobry committed 18 days ago

e455dc10

[XLA:MGPU] Port TileTransform to C++.

golechwierowicz committed 18 days ago

4abfdc62

Bumped Ruff, mypy and Pyrefly versions in .pre-commit-config.yaml

superbobry committed 18 days ago

4cb6c61b

[XLA:MGPU] Move more functions to utils and make NdIndices accept callback.

golechwierowicz committed 18 days ago

30427650

[Mosaic GPU] Fix bugs in `transform_type` implementation.

bchetioui committed 18 days ago

f9a10ff0

[XLA:MGPU] Add a convenience 'c' function constructing ConstantOp

golechwierowicz committed 18 days ago

5dacd8a4

[Pallas/Mosaic GPU] Use new API for collective metadata creation

Google-ML-Automation committed 18 days ago

55513d6f

Merge pull request #34985 from gnecula:test_pr

Google-ML-Automation committed 18 days ago

238fdfab

[NFC] Remove unnecessary comment from `testMultivariateNormalSingularCovariance`.

bchetioui committed 18 days ago

3334de9a

Add an error message to assertValidVersion

gnecula committed 18 days ago

24c453ae

Older