ROCm/jax

Pull Requests Commits

Rename top level build file to BUILD.bazel.

dfm committed 303 days ago

525cb4bd

Port many uses of contextlib.contextdecorator to explicit context manager classes.

hawkinsp committed 303 days ago

256e37af

Deprecate alpha argument to trsm LAPACK kernel.

dfm committed 303 days ago

2ce88c95

Expose pallas.mosaic.random.sample_block to pltpu interface

Google-ML-Automation committed 303 days ago

05614edc

Pass JAXLIB_* env variables to docker container

nitins17 committed 303 days ago

8ac29697

[Mosaic GPU] Use explicit recursion in rules instead of doing it automatically

apaszke committed 303 days ago

ced28167

Create `jax` wheel build target.

Google-ML-Automation committed 303 days ago

eb912ad0

Fix incorrect line separator

nitins17 committed 303 days ago

7a162f2a

Merge pull request #26668 from jakevdp:sharp-bits

Google-ML-Automation committed 303 days ago

c06c7851

Merge pull request #26734 from hawkinsp:tsan

Google-ML-Automation committed 303 days ago

69a6aaa3

Merge pull request #26564 from gspschmid:gschmid/mini_mpmd

Google-ML-Automation committed 303 days ago

a6b8384a

Typecheck pallas.CostEstimate

Google-ML-Automation committed 303 days ago

0f8e6b99

Add jax.experimental._mini_mpmd

gspschmid committed 303 days ago

2b4c455a

[Easy][Mosaic] Tiny refactor for clarity in getTypeBitwidth

Google-ML-Automation committed 303 days ago

083ffd37

[Pallas:MGPU] Adjust warpgroup lowering to the recent emit_pipeline changes

apaszke committed 303 days ago

3d87a01b

Add tests for lax.linalg.svd algorithm specification.

dfm committed 303 days ago

a3a48af1

[jax:custom_partitioning] Propagate static arguments to sharding_rule callback.

bixia1 committed 303 days ago

30348e90

[sharding_in_types] `physical_aval` should set the correct sharding on `ShapedArray` so that lowering and compilation don't crash

yashk2810 committed 303 days ago

9deb7e3d

[Mosaic GPU] Add layout inference for `scf.ForOp` and `scf.YieldOp`.

bchetioui committed 303 days ago

5024ef21

Merge pull request #26058 from gnecula:debug_info_pallas

Google-ML-Automation committed 303 days ago

7acd60c8

Update tsan suppressions.

hawkinsp committed 303 days ago

2325cf35

[Mosaic GPU] Add layout inference for `arith.Ext{F,SI,UI}Op` and `arith.Trunc{F,I}Op`.

bchetioui committed 303 days ago

5312b5e3

[Mosaic GPU][NFC] Remove unnecessary lambda wrappers from test.

bchetioui committed 303 days ago

6f966397

[better_errors] Port the Pallas debug info mechanisms to the new JAX DebugInfo.

gnecula committed 303 days ago

c4e0db6f

[Mosaic GPU] Add dialect lowering logic for splat constants.

bchetioui committed 303 days ago

5b13883f

[pallas:mosaic_gpu] Use `emit_pipeline` for pipelining in the lowering

superbobry committed 303 days ago

c13a2f95

[Pallas:MGPU] Enable lowering for .astype and scalar broadcasts

apaszke committed 303 days ago

676aceba

Skip tests using StringDType when NumPy version is below 2.0

apaszke committed 303 days ago

7fb6788d

[Pallas:MGPU] Change WG semantics convention to represent scalar arrays using scalars

apaszke committed 303 days ago

71c76220

[Pallas:MGPU] Avoid SMEM->GMEM wait if no outputs are transferred in the pipeline loop

apaszke committed 303 days ago

d0d5bba6

Older