Adds JAX-->TFjs converter

[marcvanzee]

To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.

---

This change is 

[pyu10055]

Thanks Marc, can you add a section for JAX conversion in the README file https://github.com/tensorflow/tfjs/blob/master/tfjs-converter/README.md#python-to-javascript

Reviewable status: 0 of 1 approvals obtained (waiting on @marcvanzee)

---

tfjs-converter/python/requirements.txt line 3 at r1 (raw file):

flax>=0.5.3
jax>=0.3.15
importlib_resources>=5.9.0

I am wondering if the order matters, can you put importlib_resources to the top?

[pyu10055]

to run the python tests, you need to following

cd tfjs-converter && yarn run-python-tests

Reviewable status: 0 of 1 approvals obtained (waiting on @marcvanzee)

---

tfjs-converter/python/tensorflowjs/BUILD line 53 at r1 (raw file):

    # We expect JAX to already be installed on the system, e.g. via
    # `pip install jax`.
    deps = [requirement("jax")],

add requirement("importlib_resources") to the deps list, this should remove the test error.

---

tfjs-converter/python/tensorflowjs/converters/jax_conversion.py line 22 at r1 (raw file):

from jax.experimental.jax2tf import shape_poly
import tensorflow as tf
from tensorflowjs.converters import convert_tf_saved_model

I am not sure about this, but importing from the init.py file does not seem to work with the bazel

[pyu10055]

The jax tests failed with a shape check
https://source.cloud.google.com/results/invocations/b8b3019a-5630-43b6-92c0-a0f15b2235c6/targets/%2F%2Ftfjs-converter%2Fpython%2Ftensorflowjs%2Fconverters:jax_conversion_test/tests;group=__main__.JaxConversionTest;test=test_convert_flax_bn;row=1

Reviewed 8 of 8 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @marcvanzee)

[marcvanzee]

The jax tests failed with a shape check https://source.cloud.google.com/results/invocations/b8b3019a-5630-43b6-92c0-a0f15b2235c6/targets/%2F%2Ftfjs-converter%2Fpython%2Ftensorflowjs%2Fconverters:jax_conversion_test/tests;group=__main__.JaxConversionTest;test=test_convert_flax_bn;row=1

Reviewed 8 of 8 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @marcvanzee)

Indeed, I made a recent change to jax2tf that causes average pooling to fail, see: jax-ml/jax#11804

I am currently working on a fix, and once that it in and JAX has rolled out a new version, we can proceed with this PR.

Apologies for the delay!

[marcvanzee]

@pyu10055 I fixed the bug and everything is passing now.
Note I am now installing JAX from github since the fix will only be in in the next release, and they just released a new version. Is this okay with you? Given that we are writing a blog post that we want to publish soon (1-2 weeks), it would be great if we could have this code ready. I'm happy to update the import once JAX releases a new version.

[pyu10055]

thanks I will merge this first.

Reviewed 2 of 2 files at r2.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @marcvanzee)

[pyu10055]

But we won't be able to release a new version of tfjs converter before the jax release is available.

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @marcvanzee)

[pyu10055]

the github dep is not recommended for our pip package.

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @marcvanzee)

[marcvanzee]

I will ask the JAX team when they are planning to release a new version, and if they would consider doing this soon.

[marcvanzee]

@pyu10055 after our offline discussion, I've now changed it so that we install the latest JAX release and disabled the test with a TODO. I also realized that I forgot to add content to the README, which I did as well. I also made a few minor changes to the section headers, so please take a look and make sure you are okay with it.

[pyu10055]

Reviewed 2 of 2 files at r3.
Reviewable status: complete! 2 of 1 approvals obtained (waiting on @marcvanzee)

[pyu10055]

Very excited about this new feature, thank you Marc!