ROCm/jax

Pull Requests Commits

Rename top level build file to BUILD.bazel.

dfm committed 1 year ago

525cb4bd

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

hawkinsp committed 1 year ago

256e37af

Deprecate alpha argument to trsm LAPACK kernel.

dfm committed 1 year ago

2ce88c95

Expose pallas.mosaic.random.sample_block to pltpu interface

Google-ML-Automation committed 1 year ago

05614edc

Pass JAXLIB_* env variables to docker container

nitins17 committed 1 year ago

8ac29697

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

apaszke committed 1 year ago

ced28167

Create `jax` wheel build target.

Google-ML-Automation committed 1 year ago

eb912ad0

Fix incorrect line separator

nitins17 committed 1 year ago

7a162f2a

Merge pull request #26668 from jakevdp:sharp-bits

Google-ML-Automation committed 1 year ago

c06c7851

Merge pull request #26734 from hawkinsp:tsan

Google-ML-Automation committed 1 year ago

69a6aaa3

Merge pull request #26564 from gspschmid:gschmid/mini_mpmd

Google-ML-Automation committed 1 year ago

a6b8384a

Typecheck pallas.CostEstimate

Google-ML-Automation committed 1 year ago

0f8e6b99

Add jax.experimental._mini_mpmd

gspschmid committed 1 year ago

2b4c455a

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

Google-ML-Automation committed 1 year ago

083ffd37

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

apaszke committed 1 year ago

3d87a01b

Add tests for lax.linalg.svd algorithm specification.

dfm committed 1 year ago

a3a48af1

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

bixia1 committed 1 year ago

30348e90

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

yashk2810 committed 1 year ago

9deb7e3d

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

bchetioui committed 1 year ago

5024ef21

Merge pull request #26058 from gnecula:debug_info_pallas

Google-ML-Automation committed 1 year ago

7acd60c8

Update tsan suppressions.

hawkinsp committed 1 year ago

2325cf35

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

bchetioui committed 1 year ago

5312b5e3

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

bchetioui committed 1 year ago

6f966397

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

gnecula committed 1 year ago

c4e0db6f

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

bchetioui committed 1 year ago

5b13883f

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

superbobry committed 1 year ago

c13a2f95

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

apaszke committed 1 year ago

676aceba

Skip tests using StringDType when NumPy version is below 2.0

apaszke committed 1 year ago

7fb6788d

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

apaszke committed 1 year ago

71c76220

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

apaszke committed 1 year ago

d0d5bba6

Older