ROCm/jax

Pull Requests Commits

Updated XLA commit

JnProfile committed 1 year ago

ffb07361

To avoid the inconsistency between process_index and process_id, replace backend.process_index with distributed.global_state.process_id in Jax compilation _cache_write function.

a-googler committed 1 year ago

3708336f

Merge pull request #19804 from gnecula:poly_speed2

a-googler committed 1 year ago

a23ae0d8

[shape_polyO] Performance improvements for symbolic dimension manipulations (step 2)

gnecula committed 1 year ago

eb9caf0d

Remove code that guards against int4 not being available.

hawkinsp committed 1 year ago

c55f1876

Merge pull request #19776 from gnecula:poly_split

a-googler committed 1 year ago

6a98382f

[shape_poly] Add support for jnp.split

gnecula committed 1 year ago

18698a1f

Merge pull request #19780 from gnecula:fix_jvp_astype

a-googler committed 1 year ago

ddb178ae

Fix JVP rule for convert_element_type to integer types.

gnecula committed 1 year ago

ddc24879

Update XLA dependency to use revision

a-googler committed 1 year ago

bfd29f61

Check if the Triton dialect bindings are available in lib/triton.py

superbobry committed 1 year ago

b4c8b0e4

Merge pull request #19813 from jakevdp:reuse-key

a-googler committed 1 year ago

75104375

Remove the access of `mesh.devices` from shard_map in tracing mode. It's okay to access it in eager mode

yashk2810 committed 1 year ago

da32609f

Update XLA dependency to use revision

a-googler committed 1 year ago

d8448673

[shmap] `shard_map` and `axis_index_groups` for standard collectives

ppham27 committed 1 year ago

c93f283f

Define reuse_key primitive in jax._src.prng

jakevdp committed 1 year ago

49eb7008

Fix a bug when dynamic grids are used together with scalar inputs

apaszke committed 1 year ago

b9824d7d

Inlined triton.compat

superbobry committed 1 year ago

88143624

Merge pull request #19801 from gnecula:poly_parse_neg

a-googler committed 1 year ago

547dd765

[shape_poly] Improve the symbolic expressions pretty-printer and parser.

gnecula committed 1 year ago

ed735608

Merge pull request #19777 from gnecula:poly_bug1

a-googler committed 1 year ago

5946049d

Update XLA dependency to use revision

a-googler committed 1 year ago

6a5d373a

Summarize equations in sharding mismatch errors lazily.

hawkinsp committed 1 year ago

9bb6f185

Don't compute the log message for checkpointing residuals unless it is going to be logged.

hawkinsp committed 1 year ago

7156f20b

[Pallas][Mosaic] Add support for roll op in pallas primitives.

bythew3i committed 1 year ago

ac45b8f7

Merge pull request #19796 from mattjj:jax-attrs-jit-arg

a-googler committed 1 year ago

f5ce76ba

[PJRT:C] Rename `PJRT_Structure_*` to `PJRT_Extension_*`

skye committed 1 year ago

cbe23b43

Add an interface for IFRT backends to specify executable serializability and use this information to automatically disable JAX compilation cache for backends that do not implement executable serialization

junwhanahn committed 1 year ago

b9c72967

[attrs] allow passing a jax-attrs object to jit functions

mattjj committed 1 year ago

a45cc437

[XLA:Mosaic] Expose tpu::RotateOp with stride.

bythew3i committed 1 year ago

50308553

Older