ROCm/jax

Pull Requests Commits

RNN fix MIOpen wights layout

shurale-nkn committed 39 days ago

269fa622

[internal] explicitly handle objects with dtype attribute in jax._src.dtypes.dtype

Jake VanderPlas committed 39 days ago

66a99fe9

Merge pull request #36569 from jakevdp:key-reuse-reshard

Google-ML-Automation committed 39 days ago

be39f519

Merge pull request #36521 from ahmedtaha100:complex-gamma

Google-ML-Automation committed 40 days ago

dbe4e3c1

[key reuse] add key reuse rules for sharding-related primitives

jakevdp committed 40 days ago

5ac4ab76

Add `tree_util.is_tree_node` to allow checking whether a type is a registered JAX pytree.

IvyZX committed 40 days ago

f499352c

Fix threefry2x32 to calculate the correct output sharding. Fixes https://github.com/jax-ml/jax/issues/36562

yashk2810 committed 40 days ago

8266a4d7

Add complex-input support to jax.scipy.special.gamma

ahmedtaha100 committed 40 days ago

54828727

Merge pull request #36360 from olupton:jax-toolbox-links

Google-ML-Automation committed 40 days ago

4a8d524d

[github] Upload workflow_run.json and jobs.json to GCS.

danielsuo committed 40 days ago

5745f144

Add support for parallelizing lapack invocations on the batch dimension.

hawkinsp committed 40 days ago

bfd9e1fb

Fix/remove JAX-Toolbox links.

olupton committed 40 days ago

cc420fb6

[Mosaic TPU] Allow inferring divisibility through divisions

apaszke committed 40 days ago

3ba0e1f2

[Pallas:MGPU] Make sure that the warp axis returns indices in [0, 3)

apaszke committed 40 days ago

5bd8a46e

Merge pull request #36559 from superbobry:main

Google-ML-Automation committed 40 days ago

e0696314

Upgraded to Pyrefly 0.60.0

superbobry committed 40 days ago

34bf5455

[Mosaic GPU] Enable layout inference for collective scaled MMA.

allanrenucci committed 40 days ago

03cf9812

[mGPU] Use CPU target config, that is passed from XLA side via FFI.

khasanovaa committed 40 days ago

8708d738

add a test for a simple elementwise unary kernel

ImanHosseini committed 40 days ago

6e73777c

[Mosaic GPU] Add an `AnyOf` constraint to layout inference.

allanrenucci committed 40 days ago

7141987a

Merge pull request #36549 from jakevdp:pyrefly-059

Google-ML-Automation committed 40 days ago

cbc13839

[pallas:mgpu] Allow `commit_smem_to_gmem_group` in a single warp context.

chr1sj0nes committed 40 days ago

9bf6adf3

Remove the tools/build_defs.bzl

gnecula committed 40 days ago

254515d8

Update XLA dependency to use revision http://github.com/openxla/xla/commit/1924eec184be486bce582db95964636696ea3f03

Google-ML-Automation committed 40 days ago

8552ae99

[Mosaic GPU][NFC] Add `orders_tensor_core` to DialectBarrierRef.

allanrenucci committed 40 days ago

beed3313

Merge pull request #36525 from gnecula:consts_2

Google-ML-Automation committed 40 days ago

8c1da538

Make sure vary_unreduced_cast and reduced_vary_cast order the axes wrt the context mesh. Also code-golf some code in shard_map

yashk2810 committed 40 days ago

6f55c374

Merge pull request #36551 from jakevdp:eig-autodiff-docs

Google-ML-Automation committed 40 days ago

a6285551

[pmap] Create new ifrt array when rewrapping global arrays as host local in multihost pmap.

danielsuo committed 40 days ago

6bad8fcf

Merge pull request #36535 from jakevdp:numpy-dtype-dep

Google-ML-Automation committed 40 days ago

1b046188

Older