ROCm/jax

Pull Requests Commits

Trigger CI/CD pipelinee

alekstheod committed 51 days ago

078bb39d

Trigger CI/CD pipelinee

alekstheod committed 51 days ago

3f17dfaa

Use python version

alekstheod committed 51 days ago

37d0d405

Include prebuilt rocm plugin pkgs as deps to tests

alekstheod committed 51 days ago

0fb48a15

Make wheel builds reproducible

alekstheod committed 51 days ago

949be50e

alekstheod committed 51 days ago

e2ffb999

Ignore failing tests

alekstheod committed 51 days ago

9dfb9c22

Increase timeout

alekstheod committed 51 days ago

87b2f511

Add workflow for build wheels and running bazel tests on ROCm

charleshofer committed 51 days ago

ddeb8ce6

[Pallas:MGPU] Fix a race in dynamic_scheduling_loop

apaszke committed 52 days ago

84a33c08

[typing] add type annotation for shard_map

Jake VanderPlas committed 52 days ago

709086fa

Fix a bug in how input/output memory spaces are detected and propagated through input/output aliases.

brianwa84 committed 52 days ago

9c9dd77e

Enable some tests with TC-tiling with newer libtpu.

brianwa84 committed 52 days ago

c6d2109c

[lax_numpy] Skip AdvancedIndexing tests temporarily.

danielsuo committed 52 days ago

122abbba

Merge pull request #35357 from superbobry:more-piys

Google-ML-Automation committed 52 days ago

1ae8be13

[Mosaic TPU] Allow non-DMA semaphores in EnqueueDMA

apaszke committed 52 days ago

ace7a790

[lax_numpy] Skip ReducerWhere tests temporarily.

danielsuo committed 52 days ago

dc112482

[lax_scipy] Adjust tolerance for np.float32 in trapezoid tests.

danielsuo committed 52 days ago

ab5f1971

[jaxlib] Every MLIR dialect now comes with a .pyi

superbobry committed 52 days ago

9169891b

[windows] Match triton importing in pallas_shape_poly_test to how it's imported in other tests.

danielsuo committed 52 days ago

eda5ec98

[Pallas:MGPU] Add a test for tcgen05 MMA that's both sparse and block scaled

apaszke committed 52 days ago

90fdba0f

Merge pull request #35334 from ROCm:ci_arech_fix_dangling_string_view

Google-ML-Automation committed 52 days ago

f6af222b

Add SC support for pltpu.trace_value

brianwa84 committed 52 days ago

2e3c8407

Update XLA dependency to use revision http://github.com/openxla/xla/commit/9b979f4aa0aff2bd68756d1c5d23d250213b462d

Google-ML-Automation committed 52 days ago

0f31a1b0

[Mosaic TPU] Add more complete support for 16-bit `cmpf` on TPU v5. Previously, only `arith::CmpFPredicate::ONE` was supported.

yueshengys committed 52 days ago

518b0d11

Reverts f59d8213d42d124048de9ad93792bc17152134c7

Google-ML-Automation committed 52 days ago

7b56a865

Merge pull request #35379 from mattjj:mlir-lower-fun-to-lojax

Google-ML-Automation committed 53 days ago

2cbf65c1

[NFC] Remove some unnecessary skips to have better test coverage.

yueshengys committed 53 days ago

e76e398b

[hijax] mlir.lower_fun should have lo-jaxprs

mattjj committed 53 days ago

b69d1c07

[mosaic] Verify that `tpu.memref_slice` always produces tile-aligned slices

superbobry committed 53 days ago

f59d8213

Older