ROCm/jax

Pull Requests Commits

[ROCm]: Updates hip headers path for ROCm 6.0

Rahul Batra committed 2 years ago

f997609e

[ROCm]: Dockerfile updates

Rahul Batra committed 2 years ago

b7a7f0bd

Update XLA dependency to use revision

a-googler committed 2 years ago

b512b576

Merge pull request #19437 from 8bitmp3:shard_map_title

a-googler committed 2 years ago

a1213c7d

Update XLA dependency to use revision

a-googler committed 2 years ago

0568b63b

Internal change

yashk2810 committed 2 years ago

3a0b495f

Update XLA dependency to use revision

a-googler committed 2 years ago

3b02d12f

Merge pull request #18127 from rwitten:rwitten_make_array_from_single_device_arrays_docs

a-googler committed 2 years ago

aaac4f93

Fix missing flag definition in plugin wheels built script.

Jieying Luo committed 2 years ago

b0b7c1c1

Merge pull request #19441 from jakevdp:shard-alike-fix

a-googler committed 2 years ago

f0329bf0

Added structure to make_array_from_single_device_arrays doc.

Rafi Witten committed 2 years ago

28d25a11

Return mlir modules instead of XlaComputation from custom_partitioning.

pschuh committed 2 years ago

899765ed

Guard shard_alike usage on xla_extension_version

jakevdp committed 2 years ago

80aa128e

[JAX] Disable a compilation cache test that fails on Windows in CI.

hawkinsp committed 2 years ago

a7023b18

Merge pull request #19432 from jakevdp:fix-annotation

a-googler committed 2 years ago

d74d8dd9

Add shard_map doc title

8bitmp3 committed 2 years ago

c3010ad0

[typing] fix incorrect annotation

jakevdp committed 2 years ago

3084c241

[jax2tf] Support bfloat16 for reduce_window when enable_xla=False

kevinbchen committed 2 years ago

017a0d83

[triton] Pass cluster_dims to TritonKernel and use cuLaunchKernel if size <= 1

a-googler committed 2 years ago

ab3c1b51

Add a new experimental option jax_pmap_no_rank_reduction.

hawkinsp committed 2 years ago

fc6df321

Make eval_shape a wrapper around `jax.jit(f).eval_shape(*args, **kwargs)`

yashk2810 committed 2 years ago

f04f3054

Add nested pipeline/pallas_call support for TPU meta-programming of collectives + compute.

epiqueras committed 2 years ago

85f9c51a

Updated make_array_from_single_device_arrays docs

Rafi Witten committed 2 years ago

03a8e588

Update XLA dependency to use revision

a-googler committed 2 years ago

8d29391a

Allow unhashable callables in jax.eval_shape.

mattjj committed 2 years ago

0d4f200e

[Pallas/TPU] Add support for more type conversions

sharadmv committed 2 years ago

6d21b498

[Mosaic] Fix custom call emitter handling of scalar prefetch + windowless arguments.

epiqueras committed 2 years ago

c831efb5

Merge pull request #19413 from jakevdp:dep-tie-in

a-googler committed 2 years ago

1f380e02

[Pallas/TPU] Add pallas call tests

sharadmv committed 2 years ago

3990a057

[Pallas] Use AbstractMemoryRefs for all Pallas tracing.

sharadmv committed 2 years ago

edef6d17

Older