ROCm/jax

Pull Requests Commits

add docstring for `jax.linearize` (fixes #526)

mattjj committed 6 years ago

a169e534

README should point to hosted colab for gotchas

mattjj committed 6 years ago

Verified cefbea6a

update readme, link gotchas notebook

mattjj committed 6 years ago

Verified a848d0c0

Merge pull request #488 from levskaya/master

mattjj committed 6 years ago

Verified 9711202c

fix scalar broadcast bug in safe_mul translation

mattjj committed 6 years ago

b041435f

add some simple tests for jarrett jvps

mattjj committed 6 years ago

87143d8e

handle pytrees in jarrett tangents

mattjj committed 6 years ago

122c2df2

implement manual jarrett jvps

mattjj committed 6 years ago

c9f96761

add note on dynamic slice shapes for index_update and index_add

levskaya committed 6 years ago

70686407

Merge pull request #524 from google/isclose-numpy-behavior

mattjj committed 6 years ago

Verified b1b29e52

reconcile _CheckAgainstNumpy arg order common use

mattjj committed 6 years ago

a2778b24

handle numpy < 1.14 behavior of isclose

mattjj committed 6 years ago

629c573f

Merge pull request #521 from google/remove-jaxtuples-from-optimizer-states

mattjj committed 6 years ago

Verified 35383f17

comment-out new pytree check with a TODO

mattjj committed 6 years ago

50f52070

back to original tree_mimomap optimizers solution

mattjj committed 6 years ago

f94c2063

add back an OptState that is a proper pytree

mattjj committed 6 years ago

6fdf64d7

almost done with a prefix_multimap solution

mattjj committed 6 years ago

04cfa11e

make optimizer update_fun signatures consistent

mattjj committed 6 years ago

7be36497

optimizer update must zip* the updated result too

mattjj committed 6 years ago

ce9f8bac

remove JaxTuples from optimizer states

mattjj committed 6 years ago

55a28b3a

Merge pull request #519 from google/pmap

mattjj committed 6 years ago

Verified feb9dbba

fix typos caught by review

mattjj committed 6 years ago

540eb32c

improve pmap device_num compatibility check

mattjj committed 6 years ago

e12a3cd5

avoid generating trivial lax.slice operations

mattjj committed 6 years ago

b4e83ee7

add some docstrings in pxla.py

mattjj committed 6 years ago

cb02507e

fix bug in nested pmap test code

mattjj committed 6 years ago

7129d2a0

tweak to fix xla_bridge.device_count behavior

mattjj committed 6 years ago

55512988

mattjj committed 6 years ago

ec9a9ed2

remove old jax_replica_count option

mattjj committed 6 years ago

8cfe3a38

fix typo (missing super())

mattjj committed 6 years ago

5e52abed

Newer Older