ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 364 days ago

ac1caeec

Merge pull request #29995 from mattjj:simplify-drop-unused-vars

Google-ML-Automation committed 1 year ago

11fd61f3

Merge pull request #29994 from mattjj:dedup-cleanup

Google-ML-Automation committed 1 year ago

298f4798

tweak pe._drop_unused_vars so we only call make_jaxpr_effects once

mattjj committed 1 year ago

839a0842

This change fixes 3 issues:

yashk2810 committed 1 year ago

cb44677b

split out some cleanup from #29967

mattjj committed 1 year ago

3262ee1e

Merge pull request #29676 from thisiscam:jet_unscaled

Google-ML-Automation committed 1 year ago

0df3c775

increase shard count on test file

mattjj committed 1 year ago

636a3598

implements non-factorially scaled jet computation; addresses #29624

thisiscam committed 1 year ago

354ab87e

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

ad7263c5

Reorder GSPMDSharding constructor methods to ensure fast path is taken.

tomhennigan committed 1 year ago

43b6a8d3

[Mosaic GPU] Change the definition of `EquationSystemDerivationRule` to also return a mapping from `Variable` to `VariableKey`s.

bchetioui committed 1 year ago

d0fb69f6

Add some asserts to make sure device_list is an instance of xc.DeviceList

yashk2810 committed 1 year ago

0511239d

[Pallas:MGPU] Expose TMEM_NATIVE_ROW_LAYOUT

apaszke committed 1 year ago

1838ef6a

[Mosaic GPU] Use the newly introduced scale copy API to simplify and strengthen the block scaled test

apaszke committed 1 year ago

ab81c0f6

[Mosaic GPU] Add a new API to help with copying matmul scales to TMEM

apaszke committed 1 year ago

f661ab87

[jax] Test int4 host compute

basioli-k committed 1 year ago

c5deebbe

[Pallas:MGPU] Allow explicit arrivals on Barriers with orders_tensor_core=True

apaszke committed 1 year ago

701d73a6

[Mosaic GPU][NFC] Fix the typing anotation for `TiledLayout.base_tile_shape`.

bchetioui committed 1 year ago

90b6a217

[Mosaic GPU][NFC] Fix type of `hints` in test.

bchetioui committed 1 year ago

58438387

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

06ce88e6

Merge pull request #29966 from mattjj:direct-linearize-mutable-array

Google-ML-Automation committed 1 year ago

7e7c9de8

Merge pull request #29937 from jax-ml:avoid-strong-refs

Google-ML-Automation committed 1 year ago

cd1b9520

Catch RepError in shard_map_transpose so that we raise a better error message. This is necessary because users can now annotate vma's at their level in custom_vjps.

yashk2810 committed 1 year ago

8370f433

[mutable-arrays] several fixes for direct-linearize + mutable arrays

mattjj committed 1 year ago

be053828

Merge pull request #28988 from zhenying-liu:opt_offload

Google-ML-Automation committed 1 year ago

55ff6da8

disable some tests to be broken by https://github.com/jax-ml/jax/pull/29937

mattjj committed 1 year ago

f415ad9b

If `out_shardings=None` then it means that XLA will choose the sharding for that output. But before this change under a `use_mesh(mesh)` context, we were converting `None` to replicated.

yashk2810 committed 1 year ago

44ef14da

[Mosaic GPU] Add logic to derive default layouts in an undetermined system.

bchetioui committed 1 year ago

8518505e

[Mosaic GPU] Unify `evaluate_equation` and `simplify_equation` into `reduce_equation`.

bchetioui committed 1 year ago

867c6dff

Older