ROCm/jax

Pull Requests Commits

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

bchetioui committed 298 days ago

83d89b7a

Update XLA dependency to use revision

Google-ML-Automation committed 298 days 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 298 days ago

aee89744

Merge pull request #30212 from froystig:p3f-doc

Google-ML-Automation committed 298 days ago

ae09c22f

remove caveat about random bit generation from auto-parallelization doc

froystig committed 298 days ago

60ca7bb6

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

hyeontaek committed 298 days ago

0f074062

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

yashk2810 committed 298 days ago

88b82cff

Minor cleanup in stages.py

gnecula committed 298 days ago

80205649

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

zacmustin committed 298 days ago

3ea39edf

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

superbobry committed 298 days ago

5c76ba1e

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

Google-ML-Automation committed 298 days ago

9c1efdee

[pallas] Removed some unused internal functions

superbobry committed 298 days 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 298 days ago

3f380b70

[Pallas] Fix some sus tests, make them faster

Google-ML-Automation committed 298 days ago

20404a53

Merge pull request #30167 from carlosgmartin:softmax_masked_out_clarification

Google-ML-Automation committed 298 days ago

fadfa6b9

mattjj committed 298 days ago

f927b8bc

Gather sharding rule in explicit mode

oliverdutton-iso committed 298 days 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 298 days ago

6f8456e9

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

superbobry committed 298 days ago

62ed0298

[Mosaic GPU] Add constraints to `EquationSystem`.

bchetioui committed 298 days ago

6f8ad159

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

Google-ML-Automation committed 298 days ago

d39f29ca

Update XLA dependency to use revision

Google-ML-Automation committed 299 days ago

bb6a1fba

Update XLA dependency to use revision

Google-ML-Automation committed 300 days ago

b418ee41

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

Google-ML-Automation committed 300 days ago

ff0a516f

Merge pull request #30173 from gnecula:print_large_constants

Google-ML-Automation committed 300 days ago

f56cb0ee

Summarize large constants that may appear in Jaxpr literals

gnecula committed 300 days ago

3b386790

Update XLA dependency to use revision

Google-ML-Automation committed 301 days ago

c0137e91

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

yashk2810 committed 301 days ago

440f31f6

Pallas/Mosaic support host input.

SandSnip3r committed 301 days ago

e557d4d3

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

carlosgmartin committed 301 days ago

9c1003ae

Older