Add the TensorFlow version of some JAX utilities #59

DarrenZhang01 · 2020-08-24T19:42:46Z

As titled, this Pull Request contains some necessary TF-based helper APIs for Neural Tangents, and they will be served as the main support.

folder.

tf_helpers/tf_dot_general.py

romanngg

Look great! Left a few minor comments.

tf_helpers/tf_dot_general.py

romanngg · 2020-08-24T20:19:37Z

tf_helpers/tf_dot_general.py

+Construct an equivalent general dot operation as that in JAX -
+    <https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.dot_general.html>
+
+Although there is an implementation in TF XLA, avoid directly using XLA when


Just curious, why?

Oh yeah, great question! I think Ashish mentioned that we want the newly constructed APIs - general conv, etc to be able to inter-op with the rest of the TF ecosystem, but I guess TF XLA is too low-level and independent. @wangpengmit Would you mind verifying this?

tf_helpers/tf_dot_general.py

DarrenZhang01 · 2020-08-25T01:59:43Z

Tests coming soon.

TensorFlow-related ecosystem and the tests folder contain both the `lax` and its test cases.

tf_helpers/lax.py

tf_helpers/lax_tests.py

.travis.yml

methods

…(comment)

google#59 (comment).

tf_helpers/lax.py

DarrenZhang01 · 2020-08-26T00:30:05Z

I will also add ostax into this pull request shortly.

DarrenZhang01 · 2020-08-26T00:38:05Z

Probably not going to add test cases for ostax since there are some dependencies that are still not there.

tf_helpers/tf_jax_stax.py

tf_helpers/lax.py

.travis.yml

tests/lax_test.py

tf_helpers/tf_jax_stax.py

tf_helpers/lax.py

tests/lax_test.py

the TF pool shape checker in order to make the TF `reduce_window` API consistent with JAX `reduce_window`.

stax back to TF stax.

DarrenZhang01 · 2020-08-26T20:31:59Z

Mark
#59 (comment) (Solved in bee716f),
#59 (comment) (Solved in b45f975),
https://github.com/google/neural-tangents/pull/59/files#r477001672,
https://github.com/google/neural-tangents/pull/59/files#r477558895,
https://github.com/google/neural-tangents/pull/59/files#r477003661
to do it later.

UPDATE (2nd September, 2020): Finished in #63.

tf_helpers/stax.py

tf_helpers/lax.py

google#59 (comment).

Add the TensorFlow version of general dot operation to tf_hlpers

d49ffb7

folder.

romanngg reviewed Aug 24, 2020

View reviewed changes

tf_helpers/tf_dot_general.py Outdated Show resolved Hide resolved

romanngg reviewed Aug 24, 2020

View reviewed changes

Add the TensorFlow version of some JAX lax utilities.

c57503f

DarrenZhang01 changed the title ~~Add the TensorFlow version of general dot operation to tf_hlpers~~ Add the TensorFlow version of some JAX lax utilities Aug 25, 2020

DarrenZhang01 added 4 commits August 25, 2020 15:54

Add the updated tf_lax file and the updated lax_tests file

7f51350

Remove the docstring at the begining of the file

36ae0a3

Remove unnecessary imports.

70924d0

Update Travis CI such that its installation contains the

d20c45f

TensorFlow-related ecosystem and the tests folder contain both the `lax` and its test cases.

romanngg reviewed Aug 25, 2020

View reviewed changes

DarrenZhang01 added 8 commits August 25, 2020 18:01

Rename the lax test files and move it under the test folder

2fee307

Adjust the blank lines above the class definition and between the

34787ba

methods

Remove unused comments as pointed out in google#59 (comment)

6aa0154

Remove the installation of JAX in Travis CI

42e5b54

Revert back to the original Travis CI file as mentioned in google#59 …

7dc72d6

…(comment)

Remove the _non_batched_matmul as suggested in

e0fabe3

google#59 (comment).

Remove the extra blank line in Travis CI

4dbc0b3

Give it a try on the direct file import from the tf_helpers folder

412eedc

romanngg reviewed Aug 25, 2020

View reviewed changes

tf_helpers/lax.py Outdated Show resolved Hide resolved

Remove the extra blank line and revise the import typo.

96c0ff1

Add the TF version of ostax purely for the use in Neural Tangents.

23d998d

DarrenZhang01 requested a review from romanngg August 26, 2020 00:49

DarrenZhang01 changed the title ~~Add the TensorFlow version of some JAX lax utilities~~ Add the TensorFlow version of some JAX utilities Aug 26, 2020

romanngg suggested changes Aug 26, 2020

View reviewed changes

DarrenZhang01 added 3 commits August 26, 2020 12:52

Rename tf_jax_stax to stax

c04e7b8

Remove the unused print statement

70257b9

Remove the unused tf nn import in lax tests.

d92e0fc

DarrenZhang01 added 2 commits August 26, 2020 14:48

Rename the lax tests

9432296

Remove the unused comments

51f78a4

romanngg reviewed Aug 26, 2020

View reviewed changes

tests/lax_test.py Show resolved Hide resolved

DarrenZhang01 added 2 commits August 26, 2020 16:16

Add an extra batch dimension and an extra channel dimension to pass

9357a24

the TF pool shape checker in order to make the TF `reduce_window` API consistent with JAX `reduce_window`.

Add the window_shape and strides dimension expansion that appear in JAX

27a6982

stax back to TF stax.

DarrenZhang01 added 2 commits August 26, 2020 16:38

Revert the changes in Travis CI.

a020275

Replace the vanilla NumPy support with TF NumPy support.

7df8a23

DarrenZhang01 requested review from sschoenholz and romanngg August 26, 2020 20:47

romanngg reviewed Aug 26, 2020

View reviewed changes

tf_helpers/stax.py Outdated Show resolved Hide resolved

romanngg reviewed Aug 26, 2020

View reviewed changes

tf_helpers/stax.py Outdated Show resolved Hide resolved

Remove the unused lines and unused moveaxis.

d730456

DarrenZhang01 requested a review from romanngg August 26, 2020 20:57

romanngg suggested changes Aug 26, 2020

View reviewed changes

tf_helpers/lax.py Outdated Show resolved Hide resolved

DarrenZhang01 added 2 commits August 26, 2020 17:04

Remove the extra 2 dimensions in the output of TF reduce_window.

588b748

Move np.asarray wrapper to TF pool.

ab17374

romanngg approved these changes Aug 26, 2020

View reviewed changes

romanngg merged commit a240b24 into google:neural-tangents-tf Aug 26, 2020

DarrenZhang01 pushed a commit to DarrenZhang01/neural-tangents that referenced this pull request Sep 2, 2020

Change double quotes into single quote as pointed out in

bee716f

google#59 (comment).

DarrenZhang01 pushed a commit to DarrenZhang01/neural-tangents that referenced this pull request Sep 2, 2020

Update the description of stax as pointed out in

b45f975

google#59 (comment).

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add the TensorFlow version of some JAX utilities #59

Add the TensorFlow version of some JAX utilities #59

DarrenZhang01 commented Aug 24, 2020 •

edited

romanngg left a comment

romanngg Aug 24, 2020

DarrenZhang01 Aug 24, 2020

DarrenZhang01 commented Aug 25, 2020

DarrenZhang01 commented Aug 26, 2020

DarrenZhang01 commented Aug 26, 2020

DarrenZhang01 commented Aug 26, 2020 •

edited

Add the TensorFlow version of some JAX utilities #59

Add the TensorFlow version of some JAX utilities #59

Conversation

DarrenZhang01 commented Aug 24, 2020 • edited

romanngg left a comment

Choose a reason for hiding this comment

romanngg Aug 24, 2020

Choose a reason for hiding this comment

DarrenZhang01 Aug 24, 2020

Choose a reason for hiding this comment

DarrenZhang01 commented Aug 25, 2020

DarrenZhang01 commented Aug 26, 2020

DarrenZhang01 commented Aug 26, 2020

DarrenZhang01 commented Aug 26, 2020 • edited

DarrenZhang01 commented Aug 24, 2020 •

edited

DarrenZhang01 commented Aug 26, 2020 •

edited