ROCm/jax

Pull Requests Commits

Merge pull request #19462 from superbobry:legacy-directive

a-googler committed 2 years ago

c12cf55d

DOC Add a noop implementation of the legacy directive to fix the build

superbobry committed 2 years ago

4196bf1e

Bump actions/upload-artifact from 4.1.0 to 4.2.0

dependabot[bot] committed 2 years ago

Verified 577230e2

lax.full: add sharding argument

jakevdp committed 2 years ago

f41a32c6

[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 __init__ for splash_attention

sharadmv committed 2 years ago

c03a6026

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

Newer Older