ROCm/jax

Pull Requests Commits

Reverts 6d1b5271a115007162e9f98561d6b118aa66382c

yashk2810 committed 232 days ago

8137c37e

Merge pull request #28584 from jakevdp:fix-istft

Google-ML-Automation committed 232 days ago

7a1ba3af

Merge pull request #28593 from jakevdp:fix-packbits

Google-ML-Automation committed 232 days ago

9983a7c9

Fix `with_sharding_constraint` with a scalar input

yashk2810 committed 232 days ago

4ab810b4

[Mosaic] Add explicit control over core parallelization strategy

naummo committed 232 days ago

eada2988

Rename `out_shardings -> out_sharding` in the `auto_axes` API and `in_shardings -> in_sharding` in `explicit_axes` API.

yashk2810 committed 232 days ago

5634ca12

Add basic support of `unreduced` to sharding-in-types! We cannot lower it right now, but it atleast shows up in types.

yashk2810 committed 232 days ago

48140bcb

Disable specific Triton fused_attention backwards test on A100.

justinjfu committed 232 days ago

0a68794e

Get rid of `HloModuleProto` from `CompiledMemoryStats`

junwhanahn committed 232 days ago

292b468a

[Mosaic] Move pallas + mosaic over to tpu transpose from vector. Keep vector around for compat reasons, as legacy.

Google-ML-Automation committed 232 days ago

cfab66d5

Clean up `PyLoadedExecutable::Delete`

junwhanahn committed 233 days ago

412de97e

jnp.packbits: fix handling of negative entries

jakevdp committed 233 days ago

26114161

[pallas:mosaic] Use `cf.assert` directly in the lowering rule for `checkify.check_p`

superbobry committed 233 days ago

3ea7a5da

Fold transposes feeding the RHS of a matmul into a change in the dot dimension numbers.

jaswanth- committed 233 days ago

6259cfe3

Replace `std::shared_ptr<xla::ifrt::LoadedExecutable>` with `xla::ifrt::LoadedExecutableRef`

junwhanahn committed 233 days ago

a4adf32f

expose mutable_array in experimental

Cristian Garcia committed 233 days ago

e8e586e6

Use `ArrayRef` instead of `tsl::RCReference<Array>`

junwhanahn committed 233 days ago

4cc4bd3c

Declare `tpu.vector_load`, mirroring `tpu.vector_store`.

nvgrw committed 233 days ago

63d2f7d3

Clean up `LoadedExecutable::Delete` and `LoadedExecutable::IsDeleted`

junwhanahn committed 233 days ago

59f07c3b

jax.scipy.signal.istft: support array input for window

jakevdp committed 233 days ago

f7f57d1d

[JAX] Add a test for multiprocess shard_map in McJAX with non-participating hosts. Update handling of device memories in PyDeviceList to support this.

emilyfertig committed 233 days ago

c93b051f

Merge pull request #28396 from dfm:inline-fwding

Google-ML-Automation committed 233 days ago

621fd8d0

Move const folding and forwarding into tracing

dfm committed 233 days ago

d7bc0847

[pallas:mosaic] Added a `register_lowering` decorator

superbobry committed 233 days ago

482ae677

[Mosaic GPU] Run `canonicalize` instead of `cse` before the lowering.

dimitar-asenov committed 233 days ago

ed62be8a

When registering plugins, allow to postpone the creation of the options dict.

matthiaskramm committed 233 days ago

1aee6dc6

[Mosaic GPU] Do not shortcut the transform computation for `memref.cast`

dimitar-asenov committed 233 days ago

63b51b5a

[pallas:mosaic] Fixed the type of `dimension_semantics`

superbobry committed 233 days ago

a8583db5

Exclude Mosaic GPU from the experimental target

apaszke committed 233 days ago

043b4a73

Make PartitionSpec not inherit from a tuple at runtime. For type checkers, it's still a tuple.

yashk2810 committed 233 days ago

17b02ecb

Older