ROCm/jax

Pull Requests Commits

Make lowering aware of `compute_type` so that we choose the correct lowering code.

yashk2810 committed 1 year ago

ba5b3c79

Finalize deprecation of several previously-deprecated jax.core functions:

Jake VanderPlas committed 1 year ago

fbcb157a

Be stricter when it comes to handling dtypes in splash_attention mask function

apaszke committed 1 year ago

648b9519

Merge pull request #22159 from jakevdp:arange-doc

a-googler committed 1 year ago

44071f85

[Mosaic GPU] Add support for non-128B swizzles in WGMMA

apaszke committed 1 year ago

b19ad5b3

Ported threefry2x32 for GPU to the typed XLA FFI

gnecula committed 1 year ago

cbe52429

[Mosaic GPU] Stop using nvgpu for TMA

apaszke committed 1 year ago

3ebebdfb

[Mosaic GPU] Avoid failing when importing profiler.py even if lib.mosaic_gpu is unavailable.

Google-ML-Automation committed 1 year ago

ad4c9ab8

[Mosaic] Fix bug in VectorLayout::generalizes after cl/647395486

tlongeri committed 1 year ago

3a21c81e

Disable pallas_pipeline_tpu_test under asan/msan/tsan.

apaszke committed 1 year ago

bdefb812

[Mosaic GPU] Add a warp specialized kernel with a separate TMA warp

apaszke committed 1 year ago

54edaf6e

[export] Clean up BUILD targets for jax.experimental.export

gnecula committed 1 year ago

24b42eed

Merge pull request #21976 from gnecula:exp_debugging

a-googler committed 1 year ago

016c1e88

[export] Add documentation for debugging and for ensuring compatibility.

gnecula committed 1 year ago

47f1b3de

Switches make_array_from_callback to use batched_device_put

marksandler2 committed 1 year ago

fdb1c144

Simplify mlir.dense_int_array.

hawkinsp committed 1 year ago

ac3cb6f9

Improved documentation for jnp.arange

jakevdp committed 1 year ago

3d7c53af

Add concrete layout API to JAX. The API takes `major_to_minor: tuple[int, ...]` and `tiling: tuple[tuple[int, ...], ...]` as the arguments. Allows users to pass layouts to `with_sharding_constraint` to constrain the layout + sharding.

yashk2810 committed 1 year ago

e1a496d3

Merge pull request #22157 from jakevdp:cudnn-version

a-googler committed 1 year ago

d577e299

Add a direct HLO lowering of remat_p that doesn't call eval_jaxpr.

hawkinsp committed 1 year ago

61703690

Bump minimum supported CUDNN version to 9.0

jakevdp committed 1 year ago

337cefa7

[XLA:Mosaic] Clean up pass dump and make debug easier.

bythew3i committed 1 year ago

af8bdd15

Don't enter into `JaxprEqnContext`'s manager on every equation in `jaxpr_subcomp`. Only do it where it's needed.

yashk2810 committed 1 year ago

0230d0be

Merge pull request #22147 from jakevdp:jax2tf-ci

a-googler committed 1 year ago

73176340

Merge pull request #22149 from rajasekharporeddy:matrix_rank

a-googler committed 1 year ago

d29daa00

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

55b42897

Merge pull request #22150 from dfm:xla-client-docs

a-googler committed 1 year ago

b6f73d12

Merge pull request #22129 from pkgoogle:better_divmod_doc

a-googler committed 1 year ago

f5f0f286

[Mosaic] In VectorLayout::generalizes, for (1, n) tiling, we can always squeeze out a 2nd minor dimension

tlongeri committed 1 year ago

10e598a3

Add `register_custom_call_target` to `xla_client` API docs.

dfm committed 1 year ago

dda6430f

Newer Older