ROCm/jax

Pull Requests Commits

Move ragged tests under a new class.

ghpvnist committed 1 year ago

38747a7a

[Jax][Pallas][Mosaic] Implement platform dependent diag, with branch selection driven by constant prop in mosaic lowering.

Google-ML-Automation committed 1 year ago

1719986a

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

d28b0dac

Merge pull request #25650 from jakevdp:view-int4

Google-ML-Automation committed 1 year ago

1c0dee80

Move output path to be inside the wheel build command execution loop

nitins17 committed 1 year ago

0159bead

[JAX] Add a test using inputs with different device orders for a single colocated Python call

hyeontaek committed 1 year ago

0e0fc0ac

Change the namespace name to avoid using `export` c++ keyword on namespace.

maxwillzq committed 1 year ago

e560c6a4

Accelerate deprecation of legacy JAX FFI calling convention.

dfm committed 1 year ago

4216f8fa

Merge pull request #25416 from jakevdp:api-compat

Google-ML-Automation committed 1 year ago

3a35155b

Support int4/uint4 in jnp.ndarray.view

jakevdp committed 1 year ago

75f36dc3

Merge pull request #25648 from hawkinsp:warnings3

Google-ML-Automation committed 1 year ago

44d67e13

Merge pull request #25646 from jakevdp:int4-casting

Google-ML-Automation committed 1 year ago

043c2603

Add int4/uint4 support to bitcast_convert_type

jakevdp committed 1 year ago

beee98ab

Avoid calls to warnings.catch_warnings in JAX core code.

hawkinsp committed 1 year ago

59e5ce22

[Mosaic:TPU] Roll forward of cl/708011538 (expanded trunc support), minus changes in infer-vector-layout

tlongeri committed 1 year ago

7ecc9471

Migrate _mlir Python binding target to nanobind.

hawkinsp committed 1 year ago

0ff3f144

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

aa386f85

[Mosaic GPU] Add a lowering for simple `async_load` and `async_store` ops.

dimitar-asenov committed 1 year ago

dad23fed

Merge pull request #25616 from jakevdp:abstractify

Google-ML-Automation committed 1 year ago

01e8f889

Clarify documentation for output_offsets operand of ragged_all_to_all.

olegshyshkov committed 1 year ago

db464b3f

Make jax.Arrays a necessary part of the cycle in the GC guard test

apaszke committed 1 year ago

d2f937e2

Unify abstractify & shaped_abstractify rules

jakevdp committed 1 year ago

c560f8e0

[pallas:mosaic_gpu] Change the fori tests to also take the while_p path and fix the bug.

cperivol committed 1 year ago

20efbd96

Relax some test tolerances in N-D FFT tests.

dfm committed 1 year ago

0b190bb6

Merge pull request #25625 from mattjj:ref-errors-5

Google-ML-Automation committed 1 year ago

5031b6f5

[Mosaic:TPU] Roll back cl/708011538 and cl/708112341

tlongeri committed 1 year ago

05f3a701

[Mosaic TPU] Add relayout-insertion pass and support bitwidth change for i1 vector relayout

bythew3i committed 1 year ago

2faf5402

[Mosaic:TPU] Fix trunc infer rule after cl/708011538

tlongeri committed 1 year ago

8b02884c

add mutable array ref error checks to cond and custom_vjp

mattjj committed 1 year ago

b6482f12

Merge pull request #25489 from emilyfertig:distributed-init-input-validation

Google-ML-Automation committed 1 year ago

64d3871e

Older