ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 330 days ago

e4c6cfeb

[Mosaic:TPU][Relayouts] (1, 128 * packing) <-> (n * packing, 128) retilings for n > 1

tlongeri committed 330 days ago

65e83bd8

Move jax._src.dlpack to its own build rule

Jake VanderPlas committed 330 days ago

172e3cd3

Update XLA dependency to use revision

Google-ML-Automation committed 331 days ago

bdaee513

Do not propagate name_stacks into lower_jaxpr_to_fun.

hawkinsp committed 331 days ago

ca49e4ee

Make jacrev + shmap work when check_vma=True. This involves varying the `flat_basis` output pytree on the individual output leaves instead of a union of vma of all leaves.

yashk2810 committed 331 days ago

720d4d12

Populate the HLO op_type field from JAX StableHLO.

hawkinsp committed 331 days ago

dd6f7079

Move jax._src.debugging to its own build rule

Jake VanderPlas committed 331 days ago

dc7647f3

Fix `dynamic_update_slice` transpose rule to create `zeros` on the same sharding as `update`

yashk2810 committed 332 days ago

07585ef4

Merge pull request #29817 from mattjj:mutable-array-scan-ad

Google-ML-Automation committed 332 days ago

1bd905a4

[mutable-arrays] fix autodiff with scan

mattjj committed 332 days ago

aa9ef851

Remove `multi_sharding_in_dim` and merge the operand sharding into `multi_broadcast_in_dim`. Fixes https://github.com/jax-ml/jax/issues/29818

yashk2810 committed 332 days ago

50afd1ef

Move jax._src.checkify to its own build rule.

Jake VanderPlas committed 332 days ago

303e6e01

Update XLA dependency to use revision

Google-ML-Automation committed 332 days ago

decac713

Remove JAX_IFRT_VERSION_NUMBER check for types in XLA/JAX

apivovarov committed 332 days ago

dde062dd

Enter into auto mode for `.at(...).add(...)` and set the `out_sharding` (on auto_axes) to be the same as the input sharding. Fixes https://github.com/jax-ml/jax/issues/29654

yashk2810 committed 332 days ago

f73d8c93

Add a test for bf16 1D memref load -> broadcast issue.

jaswanth- committed 332 days ago

5c9f9c04

[Mosaic] Support faster packing/unpacking for bf16 <-> f8e5m2 and f8e4m3fn.

WindQAQ committed 332 days ago

f268dfb6

Set the default of `jax_dump_ir_modes` to `stablehlo`. Fixes https://github.com/jax-ml/jax/issues/29819

yashk2810 committed 332 days ago

85150415

[Mosaic:TPU][Relayouts] (packing, 128) to (1, 128 * packing)

tlongeri committed 332 days ago

c2a21748

Add `out_sharding` to `jnp.matmul`. Fixes https://github.com/jax-ml/jax/issues/29754

yashk2810 committed 332 days ago

57ce8fb6

Include `dispatch` primitives in `roofline`.

zacmustin committed 332 days ago

fb0844ca

Add `out_sharding` to `jnp.ones`, `jnp.zeros` and `jnp.emtpy`

yashk2810 committed 332 days ago

567f567e

Allow `prng.KeyTy` as a valid `RooflineShape.dtype`.

zacmustin committed 332 days ago

d3c8de76

[Pallas MGPU] Remove unnecessary synchronization if we’re not copying out refs.

Rifur13 committed 332 days ago

92fcd21c

Reverts ebf7d6498526f601da3ffeb3c6083702ceee3940

WindQAQ committed 332 days ago

fb220c9e

Add an environment variable to enable GPU collective cancelling.

mwhittaker committed 332 days ago

e6bdd2ab

[JAX] Validate efficient resharding via `jax.device_put` with a complex mesh change

hyeontaek committed 333 days ago

747a68da

[Pallas:MGPU] Bring back plgpu.load_p to control the optimized flag

apaszke committed 333 days ago

d527ded5

Relax `rtol` in `test_shmap_unreduced_custom_vjp_bwd` for GPU backend

yashk2810 committed 333 days ago

f2a65992

Older