ROCm/jax

Pull Requests Commits

[ROCm]: Port Jax-Triton API to ROCm

Rahul Batra committed 2 years ago

6f7afead

[ROCm]: Fix duplicate deps include

Rahul Batra committed 2 years ago

8170e9ea

Fix type annotations for jnp.where

jakevdp committed 2 years ago

dd9a93b5

[JAX] Remove jax.interpreters.pxla._pmap_sharding_spec.

hawkinsp committed 2 years ago

7708cf5c

[JAX] Don't include ShardingSpecs or out_indices in the data passed to the C++ pmap() fast path.

hawkinsp committed 2 years ago

a6a8f485

[jax][benchmark] Added clearing caches for benchmarking compilation time in sparse JAX benchmarks

a-googler committed 2 years ago

f498442d

[Pallas] Import ops in jax/experimental/pallas/ops/__init__.py

laurentes committed 2 years ago

391d45fe

[Pallas] Upstream pallas to JAX

sharadmv committed 2 years ago

d872812a

Merge pull request #16843 from mattjj:shmap-partial-auto

a-googler committed 2 years ago

69cd3ebe

Add internal jaxlib function for fetching the topology from

pschuh committed 2 years ago

614bbcc6

[shmap-partial-auto] start adding partial auto / partial manual to shmap

mattjj committed 2 years ago

44da2747

Prune some exports from jax.experimental.pjit.

hawkinsp committed 2 years ago

0116d196

Enable test for indexing with u8 indices.

hawkinsp committed 2 years ago

2e042b61

Improve the repr of NamedSharding and error message of device_put

yashk2810 committed 2 years ago

853c4702

Bump minimum_jaxlib_version to 0.4.14. `xla_extension_version` is 174 and `mlir_api_version` is 54

yashk2810 committed 2 years ago

4ddf6a9a

Don't depend on jax in mesh_utils to remove circular dependency.

yashk2810 committed 2 years ago

109ed502

[Mosaic] Allow users to opt out of window prefetching

apaszke committed 2 years ago

716f4f81

Merge pull request #16900 from gnecula:poly_dot2

a-googler committed 2 years ago

b8019dcc

Merge pull request #16735 from patrick-kidger:better-debugger

a-googler committed 2 years ago

f049aee4

jax.debug.breakpoint now supports being DCE'd. Drive-by: fix #16186

patrick-kidger committed 2 years ago

c1d1be00

Consolidate the code path for both `call_tf_graph=True` or `call_tf_graph=False`.

maxwillzq committed 2 years ago

2e35f25b

[Mosaic] Add support for sublane strided store.

bythew3i committed 2 years ago

6e37c420

Fix exception if additional ABSL flags were registered after config_with_absl() was called.

hawkinsp committed 2 years ago

a40f900e

Disable testiDctn just like testiDct

yashk2810 committed 2 years ago

e1e88e1d

[Mosaic] Add support for sublane strided load.

bythew3i committed 2 years ago

9d62d867

Fix resolve_argnums if inspect.signature fails. In this case, donate_argnames was None leading an error in `assert_no_intersection`.

yashk2810 committed 2 years ago

e4955ecd

Fix the bug where args_info didn't have correct donated bit for args when donate_argnums was set on jax.jit. Fixes https://github.com/google/jax/issues/16906

yashk2810 committed 2 years ago

3929a63a

[Mosaic] Swap the order of all-reduce and elementwise ops

apaszke committed 2 years ago

5c952a1c

[shape_poly] Fix handling of dot_general with different lhs_dtype and rhs_dtype

gnecula committed 2 years ago

2eaf545a

[jax_triton] Add support for float scalar inputs.

chr1sj0nes committed 2 years ago

714156df

Older