ROCm/jax

Pull Requests Commits

[ROCm] improve gpu script

Ruturaj4 committed 1 year ago

e06be544

Merge pull request #22523 from mattjj:scan-attrs-fix

a-googler committed 1 year ago

f1974b64

fix basic scan bug with attrs

mattjj committed 1 year ago

ef82cb21

Merge pull request #23102 from rajasekharporeddy:testbranch1

a-googler committed 1 year ago

ae211a13

Merge pull request #22586 from cool-RR:logdir-pathlib

a-googler committed 1 year ago

1563897d

Normalize "interpreter mode" to "interpret mode", and "InterpreterTest" to "InterpretTest"

ayaka14732 committed 1 year ago

36739e84

Merge pull request #23045 from selamw1:diag_docstring

a-googler committed 1 year ago

d115c1cd

Better docs for jnp.fmin and fmax

rajasekharporeddy committed 1 year ago

351b780c

Merge pull request #23094 from jakevdp:basearray-annotations

a-googler committed 1 year ago

0666cccd

Merge pull request #23103 from rajasekharporeddy:testbranch2

a-googler committed 1 year ago

772e042e

Refactor FFI shape inference functions to include dimension check.

dfm committed 1 year ago

30d54ec6

Fix bug in ffi_lowering where custom layouts were ignored.

dfm committed 1 year ago

79c222ee

Rollback for: Implement initial vmap over pallas_call w/ ragged inputs (via jumbles)

apaszke committed 1 year ago

66a3f87a

Add support for shape polymorphism in ffi_lowering and move lu_pivots_to_permutation lowering out of jaxlib.

dfm committed 1 year ago

dad2f576

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

05792c95

[numpy] Fix test failures under NumPy 2.0.

hawkinsp committed 1 year ago

ba5b0815

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

b957f8ba

Remove the unused _CONV_PREFERRED_ELEMENT_TYPE_DESCRIPTION

rajasekharporeddy committed 1 year ago

e39c8685

Improve type annotations for jax.Array methods

jakevdp committed 1 year ago

dd697a9a

Add logging the jax2tf `mlir_module_serialized` module size.

maxwillzq committed 1 year ago

ce4dae1c

Match the argument name with the name in `Args` section in docstring

yashk2810 committed 1 year ago

0c543aef

Merge pull request #23088 from vfdev-5:update-build-build-py

a-googler committed 1 year ago

3a6ff86d

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

8b831b89

Minor fixes and documentation update for custom hermetic Python interpreter support.

vam-google committed 1 year ago

60cc0411

Merge pull request #23086 from ROCm:ci_build_fix

a-googler committed 1 year ago

3d942ef2

Merge pull request #23080 from jakevdp:array-doc

a-googler committed 1 year ago

1127f49d

Implement initial vmap over pallas_call w/ ragged inputs (via jumbles)

Google-ML-Automation committed 1 year ago

24394a1b

document jax.Array methods and attributes

jakevdp committed 1 year ago

e9d6fd37

Remove synchronization from GPU LU decomposition kernel by adding an async batch pointers builder.

dfm committed 1 year ago

b6306e39

Pass bazel options to requirements_update and requirements_nightly_update commands

vfdev-5 committed 1 year ago

12e8bf45

Older