ROCm/jax

Pull Requests Commits

[Pallas/Mosaic GPU] Introduce the `print_layout` primitive.

bchetioui committed 1 year ago

83d89b7a

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

f48cef57

[JAX] Use experimental DCN transfer library as a fallback for PjRt-IFRT cross-host device transfers when the PjRt plugin doesn't implement the cross-host transfer APIs.

emilyfertig committed 1 year ago

aee89744

Merge pull request #30212 from froystig:p3f-doc

Google-ML-Automation committed 1 year ago

ae09c22f

remove caveat about random bit generation from auto-parallelization doc

froystig committed 1 year ago

60ca7bb6

[JAX] Use JAX test utils to assert array equality/closeness

hyeontaek committed 1 year ago

0f074062

If operand and indices are fully replicated, allow that case to work in gather.

yashk2810 committed 1 year ago

88b82cff

Minor cleanup in stages.py

gnecula committed 1 year ago

80205649

Eliminate skiptest for C API in `memories_test` and add TPU C API config

zacmustin committed 1 year ago

3ea39edf

[pallas:mosaic] Added `"core_parallel"` to the allowed `dimension_semantics=`

superbobry committed 1 year ago

5c76ba1e

Merge pull request #30161 from mattjj:direct-linearize-fixes-6

Google-ML-Automation committed 1 year ago

9c1efdee

[pallas] Removed some unused internal functions

superbobry committed 1 year ago

b3beea46

Add `out_sharding` to `scatter_add_p` so that we can pass a value to that from gather's transpose rule thereby supporting sharding_in_types in gather transpose.

yashk2810 committed 1 year ago

3f380b70

[Pallas] Fix some sus tests, make them faster

Google-ML-Automation committed 1 year ago

20404a53

Merge pull request #30167 from carlosgmartin:softmax_masked_out_clarification

Google-ML-Automation committed 1 year ago

fadfa6b9

mattjj committed 1 year ago

f927b8bc

Gather sharding rule in explicit mode

oliverdutton-iso committed 1 year ago

4119d536

Error out if compiler_options are specified in a nested jit. They can only be specified on the top level jit.

yashk2810 committed 1 year ago

6f8456e9

[pallas:mosaic] Ensured that a nested jit in tpu_fusible_matmul_test.py has no compiler_options=

superbobry committed 1 year ago

62ed0298

[Mosaic GPU] Add constraints to `EquationSystem`.

bchetioui committed 1 year ago

6f8ad159

Merge pull request #30159 from andportnoy:aportnoy/mosaic-gpu-fix-test-remote-async-copy

Google-ML-Automation committed 1 year ago

d39f29ca

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

bb6a1fba

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

b418ee41

[Mosaic] Fold load->reshape and reshape->store into strided load/store respectively.

Google-ML-Automation committed 1 year ago

ff0a516f

Merge pull request #30173 from gnecula:print_large_constants

Google-ML-Automation committed 1 year ago

f56cb0ee

Summarize large constants that may appear in Jaxpr literals

gnecula committed 1 year ago

3b386790

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

c0137e91

Enable decorator factory pattern for `auto_axes` and `explicit_axes` API. Now users can do: `@auto_axes(...)` and `@explicit_axes(...)`

yashk2810 committed 1 year ago

440f31f6

Pallas/Mosaic support host input.

SandSnip3r committed 1 year ago

e557d4d3

Add clarification to softmax and log_softmax regarding masked-out elements.

carlosgmartin committed 1 year ago

9c1003ae

Older