Commits ROCm/jax - SemanticDiff

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

[Mosaic GPU][NFC] Import `equations` as `eqns` in `layout_inference2.py`.

bchetioui committed 1 year ago

f14c260e

Integrate LLVM at llvm/llvm-project@696c0f92e0fe

Google-ML-Automation committed 1 year ago

51e5f8cd

[mosaic] better error for in_layout size mismatch

ImanHosseini committed 1 year ago

b9f3a744

[Mosaic GPU] Propagate layouts from consumers and producers in the new layout inference.

bchetioui committed 1 year ago

2da8d573

[Pallas] Improve the error message when writing a ref to a ref

apaszke committed 1 year ago

7355f725

#sdy #mixed_serialization don't make JAX export use `SdyRoundTripExportPipeline` to stringify attributes and convert ops to StableHLO `CustomCallOp`s and back.

bartchr808 committed 1 year ago

fa04e2d3

[Pallas:MGPU] Add support for TCGEN05_TRANSPOSED layout

apaszke committed 1 year ago

aba1ca3a

[Pallas:MGPU] Add documentation for tcgen05 functions to the MGPU reference

apaszke committed 1 year ago

26af6aed

[Mosaic] Add tpu.enqueue_indirect_dma verification test on TC

naummo committed 1 year ago

aebc56d3

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

3b5bc35b

Move `jax._src.image` to its own build rule

Jake VanderPlas committed 1 year ago

ee1037e1

Merge pull request #29943 from jakevdp:core-deps

Google-ML-Automation committed 1 year ago

43ecae00

Replace old heartbeat options with simplified heartbeat options.

mwhittaker committed 1 year ago

7d412a4e

Merge pull request #29930 from jakevdp:device-doc

Google-ML-Automation committed 1 year ago

9f9cc269

Improve error message when collective APIs are called without a shard_map

yashk2810 committed 1 year ago

ae792499

[dep] remove some finalized deprecations in jax.core

jakevdp committed 1 year ago

4c28eded

Allow GSPMDSharding constructor to take in device_list (xc.DeviceList) as input along with `Sequence[jax.Device]`. This prevents extremely slow `tuple(devices) -> DeviceList` conversion in the GSPMDSharding constructor.

yashk2810 committed 1 year ago

14b2c90f

If `jax_ragged_dot_use_ragged_dot_instruction == True` then use_default_lowering() should return False since default lowering is via a pattern match.

Google-ML-Automation committed 1 year ago

28fd6007

Avoid strong refs to tracers in DynamicJaxprTrace.

dougalm committed 1 year ago

4d718ca2

Document out_sharding in jax.numpy functions

jakevdp committed 1 year ago

acc158e0