ROCm/jax

Pull Requests Commits

[ROCm]: Disable array reduce test for fp16

Rahul Batra committed 1 year ago

53e3af94

[ROCm]: More changes for jax-triton

Rahul Batra committed 1 year ago

d58cd8b8

[ROCm]: Add missing hipStreamWaitEvent API call

Rahul Batra committed 1 year ago

291029ff

[ROCm]: Add support to continue on fail

Rahul Batra committed 1 year ago

9a3d5f2a

Start the release of jax and jaxlib 0.4.25

yashk2810 committed 2 years ago

f9e20d58

Remove the deprecation of jax.tree_map for the release of 0.4.25

yashk2810 committed 2 years ago

f4045dce

Update XLA dependency to use revision

a-googler committed 2 years ago

22996bb6

[Mosaic] Fix bug in divisibility check in infer_vector_layout load and store rules

tlongeri committed 2 years ago

61aa7e89

Merge pull request #19960 from mattjj:attrs-autodiff

a-googler committed 2 years ago

072b43b2

[attrs] add linearize and vjp support

mattjj committed 2 years ago

b0b88d87

Merge pull request #19936 from tillahoffmann:fix-19935

a-googler committed 2 years ago

ae9d677b

Merge pull request #19499 from nstarman:array_api-broadcast_to-type-hint

a-googler committed 2 years ago

e1e9de0e

[Mosaic][NFC] Prefer mlir aliases for llvm classes/functions within mlir namespace for consistency

tlongeri committed 2 years ago

75cdef76

[Mosaic][apply_vector_layout][NFC] Use `LLVM_UNLIKELY` in `TPU_ASSERT_*` macros

tlongeri committed 2 years ago

8a43140c

[jax-triton] Synchronize autotuning stream with a main one

ezhulenev committed 2 years ago

3a69b807

Sort xla_flags so that the order of flags (whether added in ENV or command line) does not affect the generated cache key.

Jieying Luo committed 2 years ago

bb5997b1

Fix type annotation for array_api.broadcast_to

nstarman committed 2 years ago

b9f28572

[XLA:Mosaic] Unify ext/trunc in infer vector layout.

bythew3i committed 2 years ago

f5c00210

Add a fast path MPMD device_put test.

yashk2810 committed 2 years ago

4e61c885

[Mosaic] In apply_vector_layout, verify VectorLayout invariants that depend on target shape when loading them.

tlongeri committed 2 years ago

c9eaca22

Avoid a couple more uses of operator.attrgetter(), since it's unnecessarily slow compared to writing a function directly.

hawkinsp committed 2 years ago

b91a97d1

Avoid use of operator.attrgetter in core.find_top_Trace, since it allocates and is less efficient than a lambda.

hawkinsp committed 2 years ago

9dee375d

Increase some log levels to info in public methods.

a-googler committed 2 years ago

5413d158

Add a cache mapping complete traceback objects to MLIR IR locations.

hawkinsp committed 2 years ago

0080b36a

Reverts 051ebf04523da769f7314e08bfe6acb2732209b0

a-googler committed 2 years ago

9088e5c6

Promote `isclose` arguments to inexact dtype unless extended (fixes #19935).

tillahoffmann committed 2 years ago

2d95075e

Fix bool annotations in jax.numpy APIs

Jake VanderPlas committed 2 years ago

ac19d0f3

Add `--pre` to the libtpu nightly install instructions.

yashk2810 committed 2 years ago

4bb0a5a2

Update XLA dependency to use revision

a-googler committed 2 years ago

25416a49

[PJRT C API] Supports plugins to register a method to create topology.

Jieying Luo committed 2 years ago

e0bba8ff

Older