From b322b1c404671f041b6f771b912674c54836ef5f Mon Sep 17 00:00:00 2001
From: Fengwu Yao <fengwuyao@google.com>
Date: Thu, 10 Aug 2023 17:42:51 -0700
Subject: [PATCH 1/3] fix converter tests

---
 tfjs-converter/python/BUILD.bazel             |  15 +-
 .../python/requirements-dev_lock.txt          | 497 +++++------------
 tfjs-converter/python/requirements.txt        |  15 +-
 tfjs-converter/python/requirements_lock.txt   | 504 +++++-------------
 .../tensorflowjs/converters/converter_test.py |   4 +-
 .../converters/keras_h5_conversion_test.py    |   4 -
 .../converters/keras_tfjs_loader_test.py      |  18 +-
 7 files changed, 325 insertions(+), 732 deletions(-)

diff --git a/tfjs-converter/python/BUILD.bazel b/tfjs-converter/python/BUILD.bazel
index ca95a0a94d8..3df31fc7512 100644
--- a/tfjs-converter/python/BUILD.bazel
+++ b/tfjs-converter/python/BUILD.bazel
@@ -66,14 +66,15 @@ py_wheel(
     license = "Apache 2.0",
     python_tag = "py3",
     requires = [
-        "flax>=0.6.2,<0.6.3",
+        "flax>=0.7.2",
         "importlib_resources>=5.9.0",
-        "jax>=0.3.16",
-        "tensorflow>=2.12.0,<3",
-        "tensorflow-decision-forests>=1.3.0",
-        "six>=1.12.0,<2",
-        "tensorflow-hub>=0.13.0",
-        "packaging~=20.9",
+        "jax>=0.4.13",
+        "jaxlib>=0.4.13",
+        "tensorflow>=2.13.0,<3",
+        "tensorflow-decision-forests>=1.5.0",
+        "six>=1.16.0,<2",
+        "tensorflow-hub>=0.14.0",
+        "packaging~=23.1",
     ],
     strip_path_prefixes = [
         "tfjs-converter/python",
diff --git a/tfjs-converter/python/requirements-dev_lock.txt b/tfjs-converter/python/requirements-dev_lock.txt
index 1ea0f6b7d9c..f99ce9fcec2 100644
--- a/tfjs-converter/python/requirements-dev_lock.txt
+++ b/tfjs-converter/python/requirements-dev_lock.txt
@@ -10,6 +10,7 @@ absl-py==1.3.0 \
     # via
     #   chex
     #   optax
+    #   orbax-checkpoint
     #   tensorboard
     #   tensorflow
     #   tensorflow-decision-forests
@@ -21,6 +22,10 @@ astunparse==1.6.3 \
     --hash=sha256:5ad93a8456f0d084c3456d059fd9a92cce667963232cbf763eac3bc5b7940872 \
     --hash=sha256:c2652417f2c8b5bb325c885ae329bdf3f86424075c4fd1a128674bc6fba4b8e8
     # via tensorflow
+cached-property==1.5.2 \
+    --hash=sha256:9fa5755838eecbb2d234c3aa390bd80fbd3ac6b6869109bfc1b499f7bd89a130 \
+    --hash=sha256:df4f613cf7ad9a588cc381aaf4a512d26265ecebd5eb9e1ba12f1319eb85a6a0
+    # via orbax-checkpoint
 cachetools==5.2.0 \
     --hash=sha256:6a94c6402995a99c3970cc7e4884bb60b4a8639938157eeed436098bf9831757 \
     --hash=sha256:f9f17d2aec496a9aa6b76f53e3b614c965223c061982d434d160f930c698a9db
@@ -41,81 +46,6 @@ commonmark==0.9.1 \
     --hash=sha256:452f9dc859be7f06631ddcb328b6919c67984aca654e5fefb3914d54691aed60 \
     --hash=sha256:da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9
     # via rich
-contourpy==1.0.6 \
-    --hash=sha256:0236875c5a0784215b49d00ebbe80c5b6b5d5244b3655a36dda88105334dea17 \
-    --hash=sha256:03d1b9c6b44a9e30d554654c72be89af94fab7510b4b9f62356c64c81cec8b7d \
-    --hash=sha256:0537cc1195245bbe24f2913d1f9211b8f04eb203de9044630abd3664c6cc339c \
-    --hash=sha256:06ca79e1efbbe2df795822df2fa173d1a2b38b6e0f047a0ec7903fbca1d1847e \
-    --hash=sha256:08e8d09d96219ace6cb596506fb9b64ea5f270b2fb9121158b976d88871fcfd1 \
-    --hash=sha256:0b1e66346acfb17694d46175a0cea7d9036f12ed0c31dfe86f0f405eedde2bdd \
-    --hash=sha256:0b97454ed5b1368b66ed414c754cba15b9750ce69938fc6153679787402e4cdf \
-    --hash=sha256:0e4854cc02006ad6684ce092bdadab6f0912d131f91c2450ce6dbdea78ee3c0b \
-    --hash=sha256:12a7dc8439544ed05c6553bf026d5e8fa7fad48d63958a95d61698df0e00092b \
-    --hash=sha256:1b1ee48a130da4dd0eb8055bbab34abf3f6262957832fd575e0cab4979a15a41 \
-    --hash=sha256:1c0e1308307a75e07d1f1b5f0f56b5af84538a5e9027109a7bcf6cb47c434e72 \
-    --hash=sha256:1dedf4c64185a216c35eb488e6f433297c660321275734401760dafaeb0ad5c2 \
-    --hash=sha256:208bc904889c910d95aafcf7be9e677726df9ef71e216780170dbb7e37d118fa \
-    --hash=sha256:211dfe2bd43bf5791d23afbe23a7952e8ac8b67591d24be3638cabb648b3a6eb \
-    --hash=sha256:341330ed19074f956cb20877ad8d2ae50e458884bfa6a6df3ae28487cc76c768 \
-    --hash=sha256:344cb3badf6fc7316ad51835f56ac387bdf86c8e1b670904f18f437d70da4183 \
-    --hash=sha256:358f6364e4873f4d73360b35da30066f40387dd3c427a3e5432c6b28dd24a8fa \
-    --hash=sha256:371f6570a81dfdddbb837ba432293a63b4babb942a9eb7aaa699997adfb53278 \
-    --hash=sha256:375d81366afd547b8558c4720337218345148bc2fcffa3a9870cab82b29667f2 \
-    --hash=sha256:3a1917d3941dd58732c449c810fa7ce46cc305ce9325a11261d740118b85e6f3 \
-    --hash=sha256:4081918147fc4c29fad328d5066cfc751da100a1098398742f9f364be63803fc \
-    --hash=sha256:444fb776f58f4906d8d354eb6f6ce59d0a60f7b6a720da6c1ccb839db7c80eb9 \
-    --hash=sha256:46deb310a276cc5c1fd27958e358cce68b1e8a515fa5a574c670a504c3a3fe30 \
-    --hash=sha256:494efed2c761f0f37262815f9e3c4bb9917c5c69806abdee1d1cb6611a7174a0 \
-    --hash=sha256:50627bf76abb6ba291ad08db583161939c2c5fab38c38181b7833423ab9c7de3 \
-    --hash=sha256:5641927cc5ae66155d0c80195dc35726eae060e7defc18b7ab27600f39dd1fe7 \
-    --hash=sha256:5b117d29433fc8393b18a696d794961464e37afb34a6eeb8b2c37b5f4128a83e \
-    --hash=sha256:613c665529899b5d9fade7e5d1760111a0b011231277a0d36c49f0d3d6914bd6 \
-    --hash=sha256:6e459ebb8bb5ee4c22c19cc000174f8059981971a33ce11e17dddf6aca97a142 \
-    --hash=sha256:6f56515e7c6fae4529b731f6c117752247bef9cdad2b12fc5ddf8ca6a50965a5 \
-    --hash=sha256:730c27978a0003b47b359935478b7d63fd8386dbb2dcd36c1e8de88cbfc1e9de \
-    --hash=sha256:75a2e638042118118ab39d337da4c7908c1af74a8464cad59f19fbc5bbafec9b \
-    --hash=sha256:78ced51807ccb2f45d4ea73aca339756d75d021069604c2fccd05390dc3c28eb \
-    --hash=sha256:7ee394502026d68652c2824348a40bf50f31351a668977b51437131a90d777ea \
-    --hash=sha256:8468b40528fa1e15181cccec4198623b55dcd58306f8815a793803f51f6c474a \
-    --hash=sha256:84c593aeff7a0171f639da92cb86d24954bbb61f8a1b530f74eb750a14685832 \
-    --hash=sha256:913bac9d064cff033cf3719e855d4f1db9f1c179e0ecf3ba9fdef21c21c6a16a \
-    --hash=sha256:9447c45df407d3ecb717d837af3b70cfef432138530712263730783b3d016512 \
-    --hash=sha256:9b0e7fe7f949fb719b206548e5cde2518ffb29936afa4303d8a1c4db43dcb675 \
-    --hash=sha256:9bc407a6af672da20da74823443707e38ece8b93a04009dca25856c2d9adadb1 \
-    --hash=sha256:9e8e686a6db92a46111a1ee0ee6f7fbfae4048f0019de207149f43ac1812cf95 \
-    --hash=sha256:9fc4e7973ed0e1fe689435842a6e6b330eb7ccc696080dda9a97b1a1b78e41db \
-    --hash=sha256:a457ee72d9032e86730f62c5eeddf402e732fdf5ca8b13b41772aa8ae13a4563 \
-    --hash=sha256:a628bba09ba72e472bf7b31018b6281fd4cc903f0888049a3724afba13b6e0b8 \
-    --hash=sha256:a79d239fc22c3b8d9d3de492aa0c245533f4f4c7608e5749af866949c0f1b1b9 \
-    --hash=sha256:aa4674cf3fa2bd9c322982644967f01eed0c91bb890f624e0e0daf7a5c3383e9 \
-    --hash=sha256:acd2bd02f1a7adff3a1f33e431eb96ab6d7987b039d2946a9b39fe6fb16a1036 \
-    --hash=sha256:b3b1bd7577c530eaf9d2bc52d1a93fef50ac516a8b1062c3d1b9bcec9ebe329b \
-    --hash=sha256:b48d94386f1994db7c70c76b5808c12e23ed7a4ee13693c2fc5ab109d60243c0 \
-    --hash=sha256:b64f747e92af7da3b85631a55d68c45a2d728b4036b03cdaba4bd94bcc85bd6f \
-    --hash=sha256:b98c820608e2dca6442e786817f646d11057c09a23b68d2b3737e6dcb6e4a49b \
-    --hash=sha256:c1baa49ab9fedbf19d40d93163b7d3e735d9cd8d5efe4cce9907902a6dad391f \
-    --hash=sha256:c38c6536c2d71ca2f7e418acaf5bca30a3af7f2a2fa106083c7d738337848dbe \
-    --hash=sha256:c78bfbc1a7bff053baf7e508449d2765964d67735c909b583204e3240a2aca45 \
-    --hash=sha256:cd2bc0c8f2e8de7dd89a7f1c10b8844e291bca17d359373203ef2e6100819edd \
-    --hash=sha256:d2eff2af97ea0b61381828b1ad6cd249bbd41d280e53aea5cccd7b2b31b8225c \
-    --hash=sha256:d8834c14b8c3dd849005e06703469db9bf96ba2d66a3f88ecc539c9a8982e0ee \
-    --hash=sha256:d912f0154a20a80ea449daada904a7eb6941c83281a9fab95de50529bfc3a1da \
-    --hash=sha256:da1ef35fd79be2926ba80fbb36327463e3656c02526e9b5b4c2b366588b74d9a \
-    --hash=sha256:dbe6fe7a1166b1ddd7b6d887ea6fa8389d3f28b5ed3f73a8f40ece1fc5a3d340 \
-    --hash=sha256:dcd556c8fc37a342dd636d7eef150b1399f823a4462f8c968e11e1ebeabee769 \
-    --hash=sha256:e13b31d1b4b68db60b3b29f8e337908f328c7f05b9add4b1b5c74e0691180109 \
-    --hash=sha256:e1739496c2f0108013629aa095cc32a8c6363444361960c07493818d0dea2da4 \
-    --hash=sha256:e43255a83835a129ef98f75d13d643844d8c646b258bebd11e4a0975203e018f \
-    --hash=sha256:e626cefff8491bce356221c22af5a3ea528b0b41fbabc719c00ae233819ea0bf \
-    --hash=sha256:eadad75bf91897f922e0fb3dca1b322a58b1726a953f98c2e5f0606bd8408621 \
-    --hash=sha256:f33da6b5d19ad1bb5e7ad38bb8ba5c426d2178928bc2b2c44e8823ea0ecb6ff3 \
-    --hash=sha256:f4052a8a4926d4468416fc7d4b2a7b2a3e35f25b39f4061a7e2a3a2748c4fc48 \
-    --hash=sha256:f6ca38dd8d988eca8f07305125dec6f54ac1c518f1aaddcc14d08c01aebb6efc
-    # via matplotlib
-cycler==0.11.0 \
-    --hash=sha256:3a27e95f763a428a739d2add979fa7494c912a32c17c4c38c4d5f082cad165a3 \
-    --hash=sha256:9c87405839a19696e837b3b818fed3f5f69f16f1eec1a1ad77e043dcea9c772f
-    # via matplotlib
 debugpy==1.6.5 \
     --hash=sha256:048368f121c08b00bbded161e8583817af5055982d2722450a69efe2051621c2 \
     --hash=sha256:0f9afcc8cad6424695f3356dc9a7406d5b18e37ee2e73f34792881a44b02cc50 \
@@ -160,18 +90,18 @@ dm-tree==0.1.7 \
     --hash=sha256:d377bd621b485db42c4aeea0eabbd8f6274b89a9c338c2c1bf69a40c3b86a1fd \
     --hash=sha256:f3e2bd9b9c05d1a0039f7c128d8b055c8a05708ef569cdbbeec0a2946e425bd4
     # via chex
-flatbuffers==22.12.6 \
-    --hash=sha256:27f67c6fb102d41c911c26867bda71f1f8622176ac072fa30f668f4d023b5826 \
-    --hash=sha256:b12a3214e73f325ecd392503e821ac43a791181356f020cc3657ba829b4947c8
+etils==1.3.0 \
+    --hash=sha256:0a695ec45a982ae7c9deb437f1f251346d88b43ca59be67e961f61fe8bc8cae4 \
+    --hash=sha256:809a92ff72f12149441492cf4d9a26b56a4741dffb4dfb9c4c7b7afe055c2d28
+    # via orbax-checkpoint
+flatbuffers==23.5.26 \
+    --hash=sha256:9ea1144cac05ce5d86e2859f431c6cd5e66cd9c78c558317c7955fb8d4c78d89 \
+    --hash=sha256:c0ff356da363087b915fde4b8b45bdda73432fc17cddb3c8157472eab1422ad1
     # via tensorflow
-flax==0.6.2 \
-    --hash=sha256:9f933c87fb5762fbbf0920531e770bc385f24ef6eeb2f473641591fdbde9de89 \
-    --hash=sha256:a6247b412f14466fefcc70d043bd0facf72552d322acdda8a8700285308d390f
+flax==0.7.2 \
+    --hash=sha256:261c7b93e6d15ad80e2cedd2edb797d41b0b3c7805a54254de72a2366dc80148 \
+    --hash=sha256:7f023ece0b8b0d03019d2841dbe780e33b81ec42268bdc3e5d24c8fb0582fd7b
     # via -r tfjs-converter/python/requirements.txt
-fonttools==4.38.0 \
-    --hash=sha256:2bb244009f9bf3fa100fc3ead6aeb99febe5985fa20afbfbaa2f8946c2fbdaf1 \
-    --hash=sha256:820466f43c8be8c3009aef8b87e785014133508f0de64ec469e4efb643ae54fb
-    # via matplotlib
 gast==0.4.0 \
     --hash=sha256:40feb7b8b8434785585ab224d1568b857edb18297e5a3047f1ba012bc83b42c1 \
     --hash=sha256:b7adcdd5adbebf1adf17378da5ba3f543684dbec47b1cda1f3997e573cd542c4
@@ -269,114 +199,52 @@ idna==3.4 \
 importlib-metadata==5.1.0 \
     --hash=sha256:d5059f9f1e8e41f80e9c56c2ee58811450c31984dfa625329ffd7c0dad88a73b \
     --hash=sha256:d84d17e21670ec07990e1044a99efe8d615d860fd176fc29ef5c306068fda313
-    # via markdown
+    # via
+    #   jax
+    #   markdown
 importlib-resources==5.10.1 \
     --hash=sha256:32bb095bda29741f6ef0e5278c42df98d135391bee5f932841efc0041f748dc3 \
     --hash=sha256:c09b067d82e72c66f4f8eb12332f5efbebc9b007c0b6c40818108c9870adc363
-    # via -r tfjs-converter/python/requirements.txt
+    # via
+    #   -r tfjs-converter/python/requirements.txt
+    #   orbax-checkpoint
 isort==4.3.21 \
     --hash=sha256:54da7e92468955c4fceacd0c86bd0ec997b0e1ee80d97f67c35a78b719dccab1 \
     --hash=sha256:6e811fcb295968434526407adb8796944f1988c5b65e8139058f2014cbe100fd
     # via pylint
-jax==0.3.25 \
-    --hash=sha256:18bea69321cb95ea5ea913adfe5e2c1d453cade9d4cfd0dc814ecba9fc0cb6e3
+jax==0.4.13 \
+    --hash=sha256:03bfe6749dfe647f16f15f6616638adae6c4a7ca7167c75c21961ecfd3a3baaa
     # via
     #   -r tfjs-converter/python/requirements.txt
     #   chex
     #   flax
     #   optax
-    #   tensorflow
-jaxlib==0.3.25 \
-    --hash=sha256:09508f7000c0fa958fba29267338e8de75b31d7ea29bd79719a568c38f0f8d31 \
-    --hash=sha256:13446a8382aa9ed944c16af636ca111d0afbbead91eed5cc2dc71195045e71b3 \
-    --hash=sha256:1e59ba58c9e93c1e1cef243f2609ec0b0c0a81160c20b9555aecdea32ccd6a78 \
-    --hash=sha256:1f1448f102a9d05186f579b6931fa0c607783ecc915fdfaa482c19538affa180 \
-    --hash=sha256:2e008e0a6c10aa7e949555e98dc0471e0d550d5d7c109771e38a971b49480538 \
-    --hash=sha256:3c75c8efd3702687968820446e3fb9ff997f8a2a07ab92e33b80e2f12eab3d9a \
-    --hash=sha256:5295354ed5db111e6f3e88cdfa4010d11c33dd926ac61735b9096b4e0746aa7b \
-    --hash=sha256:6e2f4e51041b8371aa3976b5a3a9cdcdccb1bd7b040c9b1345cbf24bd28a8d19 \
-    --hash=sha256:71866aeaafbc9a12b59dcbe443353772ef235b40c53f8bd7403d39311822c276 \
-    --hash=sha256:78b29c72d0680829db9377ed9be326875849258a60b8173b4a388b34ad18bc78 \
-    --hash=sha256:9f3116389ee834b3cdeb30001b085f4b55d7741366034f041c1d377154aa5afa \
-    --hash=sha256:a50193ba0cbf879021c9d73d7bcfa7eafb9138895d057b774c301aac3701f9a5 \
-    --hash=sha256:f2d517635fd77e2729c0ab7863be0d290927d01b2abb2f5dc955c821f8b0d53e \
-    --hash=sha256:fec043cdd55f3257d02e9d8880b33860eacadcae1bd5e26f43fdd08ada23614d
+    #   orbax-checkpoint
+jaxlib==0.4.13 \
+    --hash=sha256:19ae4c316b17a49342432c69f7f89f190b975333f3f9e9e175f686a651bc7347 \
+    --hash=sha256:411334d903df07dc1ace8d52fc53c17f6bc1d55aff7f6e0e5cf61ec149f758a0 \
+    --hash=sha256:49690fcdd26560515fd15399fc3a44777e0bfc5db5c48fe76ff7bc7228e8b2fb \
+    --hash=sha256:522635d5e159401a386c79f1236c218c1f68fbb4ca6648115c3ad3c2c3f518ab \
+    --hash=sha256:532ebc4fb11386282ad63b83941d4557f4038c1144acf026f1f8565f64c7e9c0 \
+    --hash=sha256:8000c0d15c107328e8f7b7b3ac91dd822f5c287a80231882b620503ed141fa89 \
+    --hash=sha256:839173b2e9593f5e9a6d3c42852cd15070fe80a939246efbb5cf40eec815de89 \
+    --hash=sha256:a259bb35429bfbd3b76e43019dfc8f7d6ea94bb217400b78f7d0824ce07a58ac \
+    --hash=sha256:b5c0a9737efd95fe18fd7715ce30dfce476546705ea8934aad6731777a9631a5 \
+    --hash=sha256:bebb4cf001f180dc431f9604daf930c2d9cc778e4dda26f401ac939b7bac912e \
+    --hash=sha256:c230ef85712e608d0f048869766a5a63afeb2e72309943db0df9f959ab17307f \
+    --hash=sha256:d19c05c15f962e098d49b45e2758aacf19330d192ec5395f9ef136f62db90edc \
+    --hash=sha256:ea1bc9811ef7d73a15e3213115e88fe7f5d14b59d95027bea9fccc98e5a14af8 \
+    --hash=sha256:f4e9e34e5d8a6556f62fead14aee0b1614c2c6296f0078d8e6139d6aff109649 \
+    --hash=sha256:fde66a93e9be89d99e5792f677ed8e319667d6b2396865b1c52c1312844c47f9
     # via
+    #   -r tfjs-converter/python/requirements.txt
     #   chex
     #   optax
-keras==2.12.0 \
-    --hash=sha256:35c39534011e909645fb93515452e98e1a0ce23727b55d4918b9c58b2308c15e
+    #   orbax-checkpoint
+keras==2.13.1 \
+    --hash=sha256:5ce5f706f779fa7330e63632f327b75ce38144a120376b2ae1917c00fa6136af \
+    --hash=sha256:5df12cc241a015a11b65ddb452c0eeb2744fce21d9b54ba48db87492568ccc68
     # via tensorflow
-kiwisolver==1.4.4 \
-    --hash=sha256:02f79693ec433cb4b5f51694e8477ae83b3205768a6fb48ffba60549080e295b \
-    --hash=sha256:03baab2d6b4a54ddbb43bba1a3a2d1627e82d205c5cf8f4c924dc49284b87166 \
-    --hash=sha256:1041feb4cda8708ce73bb4dcb9ce1ccf49d553bf87c3954bdfa46f0c3f77252c \
-    --hash=sha256:10ee06759482c78bdb864f4109886dff7b8a56529bc1609d4f1112b93fe6423c \
-    --hash=sha256:1d1573129aa0fd901076e2bfb4275a35f5b7aa60fbfb984499d661ec950320b0 \
-    --hash=sha256:283dffbf061a4ec60391d51e6155e372a1f7a4f5b15d59c8505339454f8989e4 \
-    --hash=sha256:28bc5b299f48150b5f822ce68624e445040595a4ac3d59251703779836eceff9 \
-    --hash=sha256:2a66fdfb34e05b705620dd567f5a03f239a088d5a3f321e7b6ac3239d22aa286 \
-    --hash=sha256:2e307eb9bd99801f82789b44bb45e9f541961831c7311521b13a6c85afc09767 \
-    --hash=sha256:2e407cb4bd5a13984a6c2c0fe1845e4e41e96f183e5e5cd4d77a857d9693494c \
-    --hash=sha256:2f5e60fabb7343a836360c4f0919b8cd0d6dbf08ad2ca6b9cf90bf0c76a3c4f6 \
-    --hash=sha256:36dafec3d6d6088d34e2de6b85f9d8e2324eb734162fba59d2ba9ed7a2043d5b \
-    --hash=sha256:3fe20f63c9ecee44560d0e7f116b3a747a5d7203376abeea292ab3152334d004 \
-    --hash=sha256:41dae968a94b1ef1897cb322b39360a0812661dba7c682aa45098eb8e193dbdf \
-    --hash=sha256:4bd472dbe5e136f96a4b18f295d159d7f26fd399136f5b17b08c4e5f498cd494 \
-    --hash=sha256:4ea39b0ccc4f5d803e3337dd46bcce60b702be4d86fd0b3d7531ef10fd99a1ac \
-    --hash=sha256:5853eb494c71e267912275e5586fe281444eb5e722de4e131cddf9d442615626 \
-    --hash=sha256:5bce61af018b0cb2055e0e72e7d65290d822d3feee430b7b8203d8a855e78766 \
-    --hash=sha256:6295ecd49304dcf3bfbfa45d9a081c96509e95f4b9d0eb7ee4ec0530c4a96514 \
-    --hash=sha256:62ac9cc684da4cf1778d07a89bf5f81b35834cb96ca523d3a7fb32509380cbf6 \
-    --hash=sha256:70e7c2e7b750585569564e2e5ca9845acfaa5da56ac46df68414f29fea97be9f \
-    --hash=sha256:7577c1987baa3adc4b3c62c33bd1118c3ef5c8ddef36f0f2c950ae0b199e100d \
-    --hash=sha256:75facbe9606748f43428fc91a43edb46c7ff68889b91fa31f53b58894503a191 \
-    --hash=sha256:787518a6789009c159453da4d6b683f468ef7a65bbde796bcea803ccf191058d \
-    --hash=sha256:78d6601aed50c74e0ef02f4204da1816147a6d3fbdc8b3872d263338a9052c51 \
-    --hash=sha256:7c43e1e1206cd421cd92e6b3280d4385d41d7166b3ed577ac20444b6995a445f \
-    --hash=sha256:81e38381b782cc7e1e46c4e14cd997ee6040768101aefc8fa3c24a4cc58e98f8 \
-    --hash=sha256:841293b17ad704d70c578f1f0013c890e219952169ce8a24ebc063eecf775454 \
-    --hash=sha256:872b8ca05c40d309ed13eb2e582cab0c5a05e81e987ab9c521bf05ad1d5cf5cb \
-    --hash=sha256:877272cf6b4b7e94c9614f9b10140e198d2186363728ed0f701c6eee1baec1da \
-    --hash=sha256:8c808594c88a025d4e322d5bb549282c93c8e1ba71b790f539567932722d7bd8 \
-    --hash=sha256:8ed58b8acf29798b036d347791141767ccf65eee7f26bde03a71c944449e53de \
-    --hash=sha256:91672bacaa030f92fc2f43b620d7b337fd9a5af28b0d6ed3f77afc43c4a64b5a \
-    --hash=sha256:968f44fdbf6dd757d12920d63b566eeb4d5b395fd2d00d29d7ef00a00582aac9 \
-    --hash=sha256:9f85003f5dfa867e86d53fac6f7e6f30c045673fa27b603c397753bebadc3008 \
-    --hash=sha256:a553dadda40fef6bfa1456dc4be49b113aa92c2a9a9e8711e955618cd69622e3 \
-    --hash=sha256:a68b62a02953b9841730db7797422f983935aeefceb1679f0fc85cbfbd311c32 \
-    --hash=sha256:abbe9fa13da955feb8202e215c4018f4bb57469b1b78c7a4c5c7b93001699938 \
-    --hash=sha256:ad881edc7ccb9d65b0224f4e4d05a1e85cf62d73aab798943df6d48ab0cd79a1 \
-    --hash=sha256:b1792d939ec70abe76f5054d3f36ed5656021dcad1322d1cc996d4e54165cef9 \
-    --hash=sha256:b428ef021242344340460fa4c9185d0b1f66fbdbfecc6c63eff4b7c29fad429d \
-    --hash=sha256:b533558eae785e33e8c148a8d9921692a9fe5aa516efbdff8606e7d87b9d5824 \
-    --hash=sha256:ba59c92039ec0a66103b1d5fe588fa546373587a7d68f5c96f743c3396afc04b \
-    --hash=sha256:bc8d3bd6c72b2dd9decf16ce70e20abcb3274ba01b4e1c96031e0c4067d1e7cd \
-    --hash=sha256:bc9db8a3efb3e403e4ecc6cd9489ea2bac94244f80c78e27c31dcc00d2790ac2 \
-    --hash=sha256:bf7d9fce9bcc4752ca4a1b80aabd38f6d19009ea5cbda0e0856983cf6d0023f5 \
-    --hash=sha256:c2dbb44c3f7e6c4d3487b31037b1bdbf424d97687c1747ce4ff2895795c9bf69 \
-    --hash=sha256:c79ebe8f3676a4c6630fd3f777f3cfecf9289666c84e775a67d1d358578dc2e3 \
-    --hash=sha256:c97528e64cb9ebeff9701e7938653a9951922f2a38bd847787d4a8e498cc83ae \
-    --hash=sha256:d0611a0a2a518464c05ddd5a3a1a0e856ccc10e67079bb17f265ad19ab3c7597 \
-    --hash=sha256:d06adcfa62a4431d404c31216f0f8ac97397d799cd53800e9d3efc2fbb3cf14e \
-    --hash=sha256:d41997519fcba4a1e46eb4a2fe31bc12f0ff957b2b81bac28db24744f333e955 \
-    --hash=sha256:d5b61785a9ce44e5a4b880272baa7cf6c8f48a5180c3e81c59553ba0cb0821ca \
-    --hash=sha256:da152d8cdcab0e56e4f45eb08b9aea6455845ec83172092f09b0e077ece2cf7a \
-    --hash=sha256:da7e547706e69e45d95e116e6939488d62174e033b763ab1496b4c29b76fabea \
-    --hash=sha256:db5283d90da4174865d520e7366801a93777201e91e79bacbac6e6927cbceede \
-    --hash=sha256:db608a6757adabb32f1cfe6066e39b3706d8c3aa69bbc353a5b61edad36a5cb4 \
-    --hash=sha256:e0ea21f66820452a3f5d1655f8704a60d66ba1191359b96541eaf457710a5fc6 \
-    --hash=sha256:e7da3fec7408813a7cebc9e4ec55afed2d0fd65c4754bc376bf03498d4e92686 \
-    --hash=sha256:e92a513161077b53447160b9bd8f522edfbed4bd9759e4c18ab05d7ef7e49408 \
-    --hash=sha256:ecb1fa0db7bf4cff9dac752abb19505a233c7f16684c5826d1f11ebd9472b871 \
-    --hash=sha256:efda5fc8cc1c61e4f639b8067d118e742b812c930f708e6667a5ce0d13499e29 \
-    --hash=sha256:f0a1dbdb5ecbef0d34eb77e56fcb3e95bbd7e50835d9782a45df81cc46949750 \
-    --hash=sha256:f0a71d85ecdd570ded8ac3d1c0f480842f49a40beb423bb8014539a9f32a5897 \
-    --hash=sha256:f4f270de01dd3e129a72efad823da90cc4d6aafb64c410c9033aba70db9f1ff0 \
-    --hash=sha256:f6cb459eea32a4e2cf18ba5fcece2dbdf496384413bc1bae15583f19e567f3b2 \
-    --hash=sha256:f8ad8285b01b0d4695102546b342b493b3ccc6781fc28c8c6a1bb63e95d22f09 \
-    --hash=sha256:f9f39e2f049db33a908319cf46624a569b36983c7c78318e9726a4cb8923b26c
-    # via matplotlib
 lazy-object-proxy==1.8.0 \
     --hash=sha256:0c1c7c0433154bb7c54185714c6929acc0ba04ee1b167314a779b9025517eada \
     --hash=sha256:14010b49a2f56ec4943b6cf925f597b534ee2fe1f0738c84b3bce0c1a11ff10d \
@@ -456,53 +324,31 @@ markupsafe==2.1.1 \
     --hash=sha256:f121a1420d4e173a5d96e47e9a0c0dcff965afdf1626d28de1460815f7c4ee7a \
     --hash=sha256:fc7b548b17d238737688817ab67deebb30e8073c95749d55538ed473130ec0c7
     # via werkzeug
-matplotlib==3.6.2 \
-    --hash=sha256:0844523dfaaff566e39dbfa74e6f6dc42e92f7a365ce80929c5030b84caa563a \
-    --hash=sha256:0eda9d1b43f265da91fb9ae10d6922b5a986e2234470a524e6b18f14095b20d2 \
-    --hash=sha256:168093410b99f647ba61361b208f7b0d64dde1172b5b1796d765cd243cadb501 \
-    --hash=sha256:1836f366272b1557a613f8265db220eb8dd883202bbbabe01bad5a4eadfd0c95 \
-    --hash=sha256:19d61ee6414c44a04addbe33005ab1f87539d9f395e25afcbe9a3c50ce77c65c \
-    --hash=sha256:252957e208c23db72ca9918cb33e160c7833faebf295aaedb43f5b083832a267 \
-    --hash=sha256:32d29c8c26362169c80c5718ce367e8c64f4dd068a424e7110df1dd2ed7bd428 \
-    --hash=sha256:380d48c15ec41102a2b70858ab1dedfa33eb77b2c0982cb65a200ae67a48e9cb \
-    --hash=sha256:3964934731fd7a289a91d315919cf757f293969a4244941ab10513d2351b4e83 \
-    --hash=sha256:3cef89888a466228fc4e4b2954e740ce8e9afde7c4315fdd18caa1b8de58ca17 \
-    --hash=sha256:4426c74761790bff46e3d906c14c7aab727543293eed5a924300a952e1a3a3c1 \
-    --hash=sha256:5024b8ed83d7f8809982d095d8ab0b179bebc07616a9713f86d30cf4944acb73 \
-    --hash=sha256:52c2bdd7cd0bf9d5ccdf9c1816568fd4ccd51a4d82419cc5480f548981b47dd0 \
-    --hash=sha256:54fa9fe27f5466b86126ff38123261188bed568c1019e4716af01f97a12fe812 \
-    --hash=sha256:5ba73aa3aca35d2981e0b31230d58abb7b5d7ca104e543ae49709208d8ce706a \
-    --hash=sha256:5e16dcaecffd55b955aa5e2b8a804379789c15987e8ebd2f32f01398a81e975b \
-    --hash=sha256:5ecfc6559132116dedfc482d0ad9df8a89dc5909eebffd22f3deb684132d002f \
-    --hash=sha256:74153008bd24366cf099d1f1e83808d179d618c4e32edb0d489d526523a94d9f \
-    --hash=sha256:78ec3c3412cf277e6252764ee4acbdbec6920cc87ad65862272aaa0e24381eee \
-    --hash=sha256:795ad83940732b45d39b82571f87af0081c120feff2b12e748d96bb191169e33 \
-    --hash=sha256:7f716b6af94dc1b6b97c46401774472f0867e44595990fe80a8ba390f7a0a028 \
-    --hash=sha256:83dc89c5fd728fdb03b76f122f43b4dcee8c61f1489e232d9ad0f58020523e1c \
-    --hash=sha256:8a0ae37576ed444fe853709bdceb2be4c7df6f7acae17b8378765bd28e61b3ae \
-    --hash=sha256:8a8dbe2cb7f33ff54b16bb5c500673502a35f18ac1ed48625e997d40c922f9cc \
-    --hash=sha256:8a9d899953c722b9afd7e88dbefd8fb276c686c3116a43c577cfabf636180558 \
-    --hash=sha256:8d0068e40837c1d0df6e3abf1cdc9a34a6d2611d90e29610fa1d2455aeb4e2e5 \
-    --hash=sha256:9347cc6822f38db2b1d1ce992f375289670e595a2d1c15961aacbe0977407dfc \
-    --hash=sha256:9f335e5625feb90e323d7e3868ec337f7b9ad88b5d633f876e3b778813021dab \
-    --hash=sha256:b03fd10a1709d0101c054883b550f7c4c5e974f751e2680318759af005964990 \
-    --hash=sha256:b0ca2c60d3966dfd6608f5f8c49b8a0fcf76de6654f2eda55fc6ef038d5a6f27 \
-    --hash=sha256:b2604c6450f9dd2c42e223b1f5dca9643a23cfecc9fde4a94bb38e0d2693b136 \
-    --hash=sha256:ca0e7a658fbafcddcaefaa07ba8dae9384be2343468a8e011061791588d839fa \
-    --hash=sha256:d0e9ac04065a814d4cf2c6791a2ad563f739ae3ae830d716d54245c2b96fead6 \
-    --hash=sha256:d50e8c1e571ee39b5dfbc295c11ad65988879f68009dd281a6e1edbc2ff6c18c \
-    --hash=sha256:d840adcad7354be6f2ec28d0706528b0026e4c3934cc6566b84eac18633eab1b \
-    --hash=sha256:e0bbee6c2a5bf2a0017a9b5e397babb88f230e6f07c3cdff4a4c4bc75ed7c617 \
-    --hash=sha256:e5afe0a7ea0e3a7a257907060bee6724a6002b7eec55d0db16fd32409795f3e1 \
-    --hash=sha256:e68be81cd8c22b029924b6d0ee814c337c0e706b8d88495a617319e5dd5441c3 \
-    --hash=sha256:ec9be0f4826cdb3a3a517509dcc5f87f370251b76362051ab59e42b6b765f8c4 \
-    --hash=sha256:f04f97797df35e442ed09f529ad1235d1f1c0f30878e2fe09a2676b71a8801e0 \
-    --hash=sha256:f41e57ad63d336fe50d3a67bb8eaa26c09f6dda6a59f76777a99b8ccd8e26aec
-    # via flax
 mccabe==0.6.1 \
     --hash=sha256:ab8a6258860da4b6677da4bd2fe5dc2c659cff31b3ee4f7f5d64e79735b80d42 \
     --hash=sha256:dd8d182285a0fe56bace7f45b5e7d1a6ebcbf524e8f3bd87eb0f125271b8831f
     # via pylint
+ml-dtypes==0.2.0 \
+    --hash=sha256:022d5a4ee6be14569c2a9d1549e16f1ec87ca949681d0dca59995445d5fcdd5b \
+    --hash=sha256:1749b60348da71fd3c2ab303fdbc1965958dc50775ead41f5669c932a341cafd \
+    --hash=sha256:32107e7fa9f62db9a5281de923861325211dfff87bd23faefb27b303314635ab \
+    --hash=sha256:35b984cddbe8173b545a0e3334fe56ea1a5c3eb67c507f60d0cfde1d3fa8f8c2 \
+    --hash=sha256:36d28b8861a8931695e5a31176cad5ae85f6504906650dea5598fbec06c94606 \
+    --hash=sha256:50845af3e9a601810751b55091dee6c2562403fa1cb4e0123675cf3a4fc2c17a \
+    --hash=sha256:6488eb642acaaf08d8020f6de0a38acee7ac324c1e6e92ee0c0fea42422cb797 \
+    --hash=sha256:75015818a7fccf99a5e8ed18720cb430f3e71a8838388840f4cdf225c036c983 \
+    --hash=sha256:80d304c836d73f10605c58ccf7789c171cc229bfb678748adfb7cea2510dfd0e \
+    --hash=sha256:832a019a1b6db5c4422032ca9940a990fa104eee420f643713241b3a518977fa \
+    --hash=sha256:8faaf0897942c8253dd126662776ba45f0a5861968cf0f06d6d465f8a7bc298a \
+    --hash=sha256:bc29a0524ef5e23a7fbb8d881bdecabeb3fc1d19d9db61785d077a86cb94fab2 \
+    --hash=sha256:df6a76e1c8adf484feb138ed323f9f40a7b6c21788f120f7c78bec20ac37ee81 \
+    --hash=sha256:e70047ec2c83eaee01afdfdabee2c5b0c133804d90d0f7db4dd903360fcc537c \
+    --hash=sha256:e85ba8e24cf48d456e564688e981cf379d4c8e644db0a2f719b78de281bac2ca \
+    --hash=sha256:f00c71c8c63e03aff313bc6a7aeaac9a4f1483a921a6ffefa6d4404efd1af3d0 \
+    --hash=sha256:f08c391c2794f2aad358e6f4c70785a9a7b1df980ef4c232b3ccd4f6fe39f719
+    # via
+    #   jax
+    #   jaxlib
 msgpack==1.0.4 \
     --hash=sha256:002b5c72b6cd9b4bafd790f364b8480e859b4712e91f43014fe01e4f957b8467 \
     --hash=sha256:0a68d3ac0104e2d3510de90a1091720157c319ceeb90d74f7b5295a6bee51bae \
@@ -556,7 +402,13 @@ msgpack==1.0.4 \
     --hash=sha256:f5d869c18f030202eb412f08b28d2afeea553d6613aee89e200d7aca7ef01f5f \
     --hash=sha256:fb62ea4b62bfcb0b380d5680f9a4b3f9a2d166d9394e9bbd9666c0ee09a3645c \
     --hash=sha256:fcb8a47f43acc113e24e910399376f7277cf8508b27e5b88499f053de6b115a8
-    # via flax
+    # via
+    #   flax
+    #   orbax-checkpoint
+nest-asyncio==1.5.7 \
+    --hash=sha256:5301c82941b550b3123a1ea772ba9a1c80bad3a182be8c1a5ae6ad3be57a9657 \
+    --hash=sha256:6a80f7b98f24d9083ed24608977c09dd608d83f91cccc24c9d2cba6d10e01c10
+    # via orbax-checkpoint
 numpy==1.23.5 \
     --hash=sha256:01dd17cbb340bf0fc23981e52e1d18a9d4050792e8fb8363cecbf066a84b827d \
     --hash=sha256:06005a2ef6014e9956c09ba07654f9837d9e26696a0470e42beedadb78c11b07 \
@@ -588,14 +440,14 @@ numpy==1.23.5 \
     --hash=sha256:f9a909a8bae284d46bbfdefbdd4a262ba19d3bc9921b1e76126b1d21c3c34135
     # via
     #   chex
-    #   contourpy
     #   flax
     #   h5py
     #   jax
     #   jaxlib
-    #   matplotlib
+    #   ml-dtypes
     #   opt-einsum
     #   optax
+    #   orbax-checkpoint
     #   pandas
     #   scipy
     #   tensorboard
@@ -617,12 +469,15 @@ optax==0.1.4 \
     --hash=sha256:12fcf33bd682f9a162a3deb097f864130c3224d76771af2ba09410de80399a9b \
     --hash=sha256:fb7a0550d57a6636164a3de25986a8a19be8ff6431fcdf1225b4e05175810f22
     # via flax
-packaging==20.9 \
-    --hash=sha256:5b327ac1320dc863dca72f4514ecc086f31186744b84a230374cc1fd776feae5 \
-    --hash=sha256:67714da7f7bc052e064859c05c595155bd1ee9f69f76557e21f051443c20947a
+orbax-checkpoint==0.2.3 \
+    --hash=sha256:155e0a2dceef2901122e66585171e1dff4f4a4d9d2abe43a2b514279b9a3dabd \
+    --hash=sha256:a001bf48f1cebc635b07263fa546473ea48be3e278c50d5ade880b9aafb96f8a
+    # via flax
+packaging==23.1 \
+    --hash=sha256:994793af429502c4ea2ebf6bf664629d07c1a9fe974af92966e4b8d2df7edc61 \
+    --hash=sha256:a392980d2b6cffa644431898be54b0045151319d1e7ec34f0cfed48767dd334f
     # via
     #   -r tfjs-converter/python/requirements.txt
-    #   matplotlib
     #   tensorflow
 pandas==1.5.2 \
     --hash=sha256:0183cb04a057cc38fde5244909fca9826d5d57c4a5b7390c0cc3fa7acd9fa883 \
@@ -653,69 +508,6 @@ pandas==1.5.2 \
     --hash=sha256:e9dbacd22555c2d47f262ef96bb4e30880e5956169741400af8b306bbb24a273 \
     --hash=sha256:f6257b314fc14958f8122779e5a1557517b0f8e500cfb2bd53fa1f75a8ad0af2
     # via tensorflow-decision-forests
-pillow==9.3.0 \
-    --hash=sha256:03150abd92771742d4a8cd6f2fa6246d847dcd2e332a18d0c15cc75bf6703040 \
-    --hash=sha256:073adb2ae23431d3b9bcbcff3fe698b62ed47211d0716b067385538a1b0f28b8 \
-    --hash=sha256:0b07fffc13f474264c336298d1b4ce01d9c5a011415b79d4ee5527bb69ae6f65 \
-    --hash=sha256:0b7257127d646ff8676ec8a15520013a698d1fdc48bc2a79ba4e53df792526f2 \
-    --hash=sha256:12ce4932caf2ddf3e41d17fc9c02d67126935a44b86df6a206cf0d7161548627 \
-    --hash=sha256:15c42fb9dea42465dfd902fb0ecf584b8848ceb28b41ee2b58f866411be33f07 \
-    --hash=sha256:18498994b29e1cf86d505edcb7edbe814d133d2232d256db8c7a8ceb34d18cef \
-    --hash=sha256:1c7c8ae3864846fc95f4611c78129301e203aaa2af813b703c55d10cc1628535 \
-    --hash=sha256:22b012ea2d065fd163ca096f4e37e47cd8b59cf4b0fd47bfca6abb93df70b34c \
-    --hash=sha256:276a5ca930c913f714e372b2591a22c4bd3b81a418c0f6635ba832daec1cbcfc \
-    --hash=sha256:2e0918e03aa0c72ea56edbb00d4d664294815aa11291a11504a377ea018330d3 \
-    --hash=sha256:3033fbe1feb1b59394615a1cafaee85e49d01b51d54de0cbf6aa8e64182518a1 \
-    --hash=sha256:3168434d303babf495d4ba58fc22d6604f6e2afb97adc6a423e917dab828939c \
-    --hash=sha256:32a44128c4bdca7f31de5be641187367fe2a450ad83b833ef78910397db491aa \
-    --hash=sha256:3dd6caf940756101205dffc5367babf288a30043d35f80936f9bfb37f8355b32 \
-    --hash=sha256:40e1ce476a7804b0fb74bcfa80b0a2206ea6a882938eaba917f7a0f004b42502 \
-    --hash=sha256:41e0051336807468be450d52b8edd12ac60bebaa97fe10c8b660f116e50b30e4 \
-    --hash=sha256:4390e9ce199fc1951fcfa65795f239a8a4944117b5935a9317fb320e7767b40f \
-    --hash=sha256:502526a2cbfa431d9fc2a079bdd9061a2397b842bb6bc4239bb176da00993812 \
-    --hash=sha256:51e0e543a33ed92db9f5ef69a0356e0b1a7a6b6a71b80df99f1d181ae5875636 \
-    --hash=sha256:57751894f6618fd4308ed8e0c36c333e2f5469744c34729a27532b3db106ee20 \
-    --hash=sha256:5d77adcd56a42d00cc1be30843d3426aa4e660cab4a61021dc84467123f7a00c \
-    --hash=sha256:655a83b0058ba47c7c52e4e2df5ecf484c1b0b0349805896dd350cbc416bdd91 \
-    --hash=sha256:68943d632f1f9e3dce98908e873b3a090f6cba1cbb1b892a9e8d97c938871fbe \
-    --hash=sha256:6c738585d7a9961d8c2821a1eb3dcb978d14e238be3d70f0a706f7fa9316946b \
-    --hash=sha256:73bd195e43f3fadecfc50c682f5055ec32ee2c933243cafbfdec69ab1aa87cad \
-    --hash=sha256:772a91fc0e03eaf922c63badeca75e91baa80fe2f5f87bdaed4280662aad25c9 \
-    --hash=sha256:77ec3e7be99629898c9a6d24a09de089fa5356ee408cdffffe62d67bb75fdd72 \
-    --hash=sha256:7db8b751ad307d7cf238f02101e8e36a128a6cb199326e867d1398067381bff4 \
-    --hash=sha256:801ec82e4188e935c7f5e22e006d01611d6b41661bba9fe45b60e7ac1a8f84de \
-    --hash=sha256:82409ffe29d70fd733ff3c1025a602abb3e67405d41b9403b00b01debc4c9a29 \
-    --hash=sha256:828989c45c245518065a110434246c44a56a8b2b2f6347d1409c787e6e4651ee \
-    --hash=sha256:829f97c8e258593b9daa80638aee3789b7df9da5cf1336035016d76f03b8860c \
-    --hash=sha256:871b72c3643e516db4ecf20efe735deb27fe30ca17800e661d769faab45a18d7 \
-    --hash=sha256:89dca0ce00a2b49024df6325925555d406b14aa3efc2f752dbb5940c52c56b11 \
-    --hash=sha256:90fb88843d3902fe7c9586d439d1e8c05258f41da473952aa8b328d8b907498c \
-    --hash=sha256:97aabc5c50312afa5e0a2b07c17d4ac5e865b250986f8afe2b02d772567a380c \
-    --hash=sha256:9aaa107275d8527e9d6e7670b64aabaaa36e5b6bd71a1015ddd21da0d4e06448 \
-    --hash=sha256:9f47eabcd2ded7698106b05c2c338672d16a6f2a485e74481f524e2a23c2794b \
-    --hash=sha256:a0a06a052c5f37b4ed81c613a455a81f9a3a69429b4fd7bb913c3fa98abefc20 \
-    --hash=sha256:ab388aaa3f6ce52ac1cb8e122c4bd46657c15905904b3120a6248b5b8b0bc228 \
-    --hash=sha256:ad58d27a5b0262c0c19b47d54c5802db9b34d38bbf886665b626aff83c74bacd \
-    --hash=sha256:ae5331c23ce118c53b172fa64a4c037eb83c9165aba3a7ba9ddd3ec9fa64a699 \
-    --hash=sha256:af0372acb5d3598f36ec0914deed2a63f6bcdb7b606da04dc19a88d31bf0c05b \
-    --hash=sha256:afa4107d1b306cdf8953edde0534562607fe8811b6c4d9a486298ad31de733b2 \
-    --hash=sha256:b03ae6f1a1878233ac620c98f3459f79fd77c7e3c2b20d460284e1fb370557d4 \
-    --hash=sha256:b0915e734b33a474d76c28e07292f196cdf2a590a0d25bcc06e64e545f2d146c \
-    --hash=sha256:b4012d06c846dc2b80651b120e2cdd787b013deb39c09f407727ba90015c684f \
-    --hash=sha256:b472b5ea442148d1c3e2209f20f1e0bb0eb556538690fa70b5e1f79fa0ba8dc2 \
-    --hash=sha256:b59430236b8e58840a0dfb4099a0e8717ffb779c952426a69ae435ca1f57210c \
-    --hash=sha256:b90f7616ea170e92820775ed47e136208e04c967271c9ef615b6fbd08d9af0e3 \
-    --hash=sha256:b9a65733d103311331875c1dca05cb4606997fd33d6acfed695b1232ba1df193 \
-    --hash=sha256:bac18ab8d2d1e6b4ce25e3424f709aceef668347db8637c2296bcf41acb7cf48 \
-    --hash=sha256:bca31dd6014cb8b0b2db1e46081b0ca7d936f856da3b39744aef499db5d84d02 \
-    --hash=sha256:be55f8457cd1eac957af0c3f5ece7bc3f033f89b114ef30f710882717670b2a8 \
-    --hash=sha256:c7025dce65566eb6e89f56c9509d4f628fddcedb131d9465cacd3d8bac337e7e \
-    --hash=sha256:c935a22a557a560108d780f9a0fc426dd7459940dc54faa49d83249c8d3e760f \
-    --hash=sha256:dbb8e7f2abee51cef77673be97760abff1674ed32847ce04b4af90f610144c7b \
-    --hash=sha256:e6ea6b856a74d560d9326c0f5895ef8050126acfdc7ca08ad703eb0081e82b74 \
-    --hash=sha256:ebf2029c1f464c59b8bdbe5143c79fa2045a581ac53679733d3a91d400ff9efb \
-    --hash=sha256:f1ff2ee69f10f13a9596480335f406dd1f70c3650349e2be67ca3139280cade0
-    # via matplotlib
 prompt-toolkit==1.0.14 \
     --hash=sha256:7281b5199235adaef6980942840c43753e4ab20dfe41338da634fb41c194f9d8 \
     --hash=sha256:82c7f8e07d7a0411ff5367a5a8ff520f0112b9179f3e599ee8ad2ad9b943d911 \
@@ -762,18 +554,10 @@ pylint==2.5.0 ; python_version > "3.0" \
     --hash=sha256:588e114e3f9a1630428c35b7dd1c82c1c93e1b0e78ee312ae4724c5e1a1e0245 \
     --hash=sha256:bd556ba95a4cf55a1fc0004c00cf4560b1e70598a54a74c6904d933c8f3bd5a8
     # via -r tfjs-converter/python/requirements-dev.txt
-pyparsing==3.0.9 \
-    --hash=sha256:2b020ecf7d21b687f219b71ecad3631f644a47f01403fa1d1036b0c6416d70fb \
-    --hash=sha256:5026bae9a10eeaefb61dab2f09052b9f4307d44aee4eda64b309723d8d206bbc
-    # via
-    #   matplotlib
-    #   packaging
 python-dateutil==2.8.2 \
     --hash=sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86 \
     --hash=sha256:961d03dc3453ebbc59dbdea9e4e11c5651520a876d0f4db161e8674aae935da9
-    # via
-    #   matplotlib
-    #   pandas
+    # via pandas
 pytz==2022.6 \
     --hash=sha256:222439474e9c98fced559f1709d89e6c9cbf8d79c794ff3eb9f8800064291427 \
     --hash=sha256:e89512406b793ca39f5971bc999cc538ce125c0e51c27941bef4568b460095e2
@@ -819,7 +603,9 @@ pyyaml==6.0 \
     --hash=sha256:dbad0e9d368bb989f4515da330b88a057617d16b6a8245084f1b05400f24609f \
     --hash=sha256:e61ceaab6f49fb8bdfaa0f92c4b57bcfbea54c09277b1b4f7ac376bfb7a7c174 \
     --hash=sha256:f84fbc98b019fef2ee9a1cb3ce93e3187a6df0b2538a651bfb890254ba9f90b5
-    # via flax
+    # via
+    #   flax
+    #   orbax-checkpoint
 regex==2022.10.31 \
     --hash=sha256:052b670fafbe30966bbe5d025e90b2a491f85dfe5b2583a163b5e60a85a321ad \
     --hash=sha256:0653d012b3bf45f194e5e6a41df9258811ac8fc395579fa82958a8b76286bea4 \
@@ -972,47 +758,57 @@ six==1.16.0 \
     #   python-dateutil
     #   tensorflow
     #   tensorflow-decision-forests
-tensorboard==2.12.2 \
-    --hash=sha256:811ab0d27a139445836db9fd4f974424602c3dce12379364d379bcba7c783a68
+tensorboard==2.13.0 \
+    --hash=sha256:ab69961ebddbddc83f5fa2ff9233572bdad5b883778c35e4fe94bf1798bd8481
     # via tensorflow
 tensorboard-data-server==0.7.0 \
     --hash=sha256:64aa1be7c23e80b1a42c13b686eb0875bb70f5e755f4d2b8de5c1d880cf2267f \
     --hash=sha256:753d4214799b31da7b6d93837959abebbc6afa86e69eacf1e9a317a48daa31eb \
     --hash=sha256:eb7fa518737944dbf4f0cf83c2e40a7ac346bf91be2e6a0215de98be74e85454
     # via tensorboard
-tensorboard-plugin-wit==1.8.1 \
-    --hash=sha256:ff26bdd583d155aa951ee3b152b3d0cffae8005dc697f72b44a8e8c2a77a8cbe
-    # via tensorboard
-tensorflow==2.12.0 \
-    --hash=sha256:020d6a54cb26020bdc71a7bae8ee35be05096f63e773dc517f6e87c49de62c50 \
-    --hash=sha256:23850332f1f9f778d697c9dba63ca52be72cb73363e75ad358f07ddafef63c01 \
-    --hash=sha256:31f81eb8adaeb558963f5d8b47dbfcc398d898f0857bf3de6b6484350236b7b5 \
-    --hash=sha256:357d9d2851188a8d27ee195345b4d175cad970150d1344ba9d9fcc4bf2b68336 \
-    --hash=sha256:42fc2635e9420faee781a16bd393126f29cd39aa2b9d02901f24d8497bd6f958 \
-    --hash=sha256:4afc2dd57435f29ebe249eb5f595d89b0e73be94922eeb7110aa6280a332837c \
-    --hash=sha256:6e7641e2a6e32f31ff233495478a9cc86b7c038140eab714a61eeddbbbb327c3 \
-    --hash=sha256:6ec4a2934ea19e92f27a9668ece43025ed5efe14b5d19be53b07692bc8a4189d \
-    --hash=sha256:76414355e420edb9154b4e72113eef5813ccb71701fda959afbbc1eebe3099bd \
-    --hash=sha256:91dccda42c03569d8c787190482a11ecae3b9b173aaa9166f0ab20cecc9c31f4 \
-    --hash=sha256:9f70a8f9ab46e5ed436850aa60d1cd40645f5c669e14bcad48915dc1f597dda2 \
-    --hash=sha256:a7194e744c5a7f3e759ecb949527b4a07718a6d1110e6e82fd4ce0c5586a7d4a \
-    --hash=sha256:be4ac0dfcc7a16f6df2bc19bd322e312235ab3f7b0c7297f96c92c44bb14d2a1 \
-    --hash=sha256:c5193ddb3bb5120cb445279beb08ed9e74a85a4eeb2485550d6fb707a89d9a88 \
-    --hash=sha256:c8001210df7202ef6267150865b0b79f834c3ca69ee3132277de8eeb994dffde \
-    --hash=sha256:e29fcf6cfd069aefb4b44f357cccbb4415a5a3d7b5b516eaf4450062fe40021e
+tensorflow==2.13.0 \
+    --hash=sha256:00060c5516a61e30c51936084ebc37091d116efe9ae74b2818cbd8b2006218e7 \
+    --hash=sha256:06559eeaa69e6561cccbe2d02b015bcec663e875c8bbc4643f55692476e52147 \
+    --hash=sha256:076d953a1508dc58bf95f30f58bcc9ee364b1353c61e143cb20c2dada91afb05 \
+    --hash=sha256:11ad6a7ff49b4a690fbf37a5adaf28ba2686350a859c5f13c58dc8d2cc670375 \
+    --hash=sha256:19ee67901702b26787ad685cca65730c163c101c0c2f238a2584d714e0fa8c25 \
+    --hash=sha256:2822ac48c38d69b7fb104e606dacbd763c4bf5d3b20791f25be16a5076d01912 \
+    --hash=sha256:5e0fdadec59de3d11c5b5129ddc38e739bde7aa13095b82e19d4380e14d04999 \
+    --hash=sha256:6fff426661d286a4c634da44275d2ea2b951f392f3e65c8603681e7cb040586a \
+    --hash=sha256:72d68b8c2f382e2d01b956c8ba516c0a7d5dad98111dd351bf82bfa646aa1c72 \
+    --hash=sha256:7a08c0e2938ed5b642a8787678123827477b81d316055d5073fff82fa183eb82 \
+    --hash=sha256:89125443e998548059c4e4a129dfab2b1ff7f2fd4c8eaed97842c3cd9b663101 \
+    --hash=sha256:948003b5a23b72b3d89746d729e62ec5f01e47460f05521b2211d95069f569ba \
+    --hash=sha256:9c04bc3023b6c4cfb9ee9759c3f03f21993891b4c345df52eb5519204fbf28c0 \
+    --hash=sha256:b2978b39e8b3919059b5fd9e28508d50a77965d06ed0b537ed71c97de22dabdf \
+    --hash=sha256:cbb83561bb7d55859eaefc70c674e58713d4e10c10927423ed836a5289bbfa86 \
+    --hash=sha256:de77306c0c22c9d8754f54700752ac3a1efee895c5357308e6594436404bfbc0 \
+    --hash=sha256:e0cf94d36ceaba8f158c6e15404a81fd5b3aa4cb04147c674cf55bd1aec78154 \
+    --hash=sha256:e8f0b69ee2f800399fc6bc7ec55fecfa33662d136e425485959d90638f32a32a \
+    --hash=sha256:fa7abe265cc3ebccc9b405a280bf674824c6d85df5e6ccfa985987b3c9d265b4 \
+    --hash=sha256:fb2ff1129c93e853c19897d6a22ed0ec56387f5c6290ec03dec1c6f7b80bc396
     # via
     #   -r tfjs-converter/python/requirements.txt
     #   tensorflow-decision-forests
-tensorflow-decision-forests==1.3.0 \
-    --hash=sha256:9729973a6b16ee75161cf2e8d68a8fddef3988db8e05734b94a3fad6780a916d \
-    --hash=sha256:bb8b07d3e39e00153428288202e1393991308411c0af542afdc7b721856aa15a \
-    --hash=sha256:dbd9b8563af9ad5d56381fe483d75e593c8ecb820f1f671a79cc74fafc906225
+tensorflow-decision-forests==1.5.0 \
+    --hash=sha256:04fd913627d08fe54514b179c612e87eebf55f1448bf01951660985dfa14a6e1 \
+    --hash=sha256:1209a2832ac65f8f74bd9d0c1d58f3f8b771e7fa5c9d504c547842311647b7d4 \
+    --hash=sha256:22e3835acbfbd5356bb2f8e0c973dfc40ef80a7924b793e90c811158448cfe77 \
+    --hash=sha256:43ffd4fba1c3376f58a9dcee943df80f0cff6e47224d109ad0389a723c74947c \
+    --hash=sha256:4a0df3a3be5751594d49f5a8f99977b553cf1c42f320b952ac2a2f67b85283f5 \
+    --hash=sha256:5ec4297eb5e7c4110cf8aae89e9b08b9ad2cb725e3e63c89c78304c0d7235d24 \
+    --hash=sha256:804f6bed277b5c5b6d2bd85738a64973d5d3e8e6ac06abf6098545740245cedc \
+    --hash=sha256:a43af2a5a8c34e550bf549c6cad96da271979efc5a8ec988f6f76cc90770415a \
+    --hash=sha256:d137241dad8e884d0c937aa8769fe0768324804e9ba666a78b7b5f2f536a0bd2 \
+    --hash=sha256:d685e92abe44920ee6d89394ec4e075bb1ada7402f673566146e1b476a576e96 \
+    --hash=sha256:f5d8c3730578bda55a8f520ae39b0c9b2560d69bd53b57882e5371c1a82ba098 \
+    --hash=sha256:fbd403acf736bb9b4afd2985d9056e6d5043fc4b9a31bd05e5fcae2b1d413dc3
     # via -r tfjs-converter/python/requirements.txt
-tensorflow-estimator==2.12.0 \
-    --hash=sha256:59b191bead4883822de3d63ac02ace11a83bfe6c10d64d0c4dfde75a50e60ca1
+tensorflow-estimator==2.13.0 \
+    --hash=sha256:6f868284eaa654ae3aa7cacdbef2175d0909df9fcf11374f5166f8bf475952aa
     # via tensorflow
-tensorflow-hub==0.13.0 ; python_version >= "3" \
-    --hash=sha256:3544f4fd9fd99e4eeb6da1b5b5320e4a2dbdef7f9bb778f66f76d6790f32dd65
+tensorflow-hub==0.14.0 \
+    --hash=sha256:519c6b56c4d304667fbd8ce66bd637e6a750c901215468db2cc6bfd0739bb0b0
     # via -r tfjs-converter/python/requirements.txt
 tensorflow-io-gcs-filesystem==0.28.0 \
     --hash=sha256:00cf6a92f1f9f90b2ba2d728870bcd2a70b116316d0817ab0b91dd390c25b3fd \
@@ -1031,24 +827,27 @@ tensorflow-io-gcs-filesystem==0.28.0 \
     --hash=sha256:cc062ce13ec95fb64b1fd426818a6d2b0e5be9692bc0e43a19cce115b6da4336 \
     --hash=sha256:f76cbe1a784841c223f6861e5f6c7e53aa6232cb626d57e76881a0638c365de6
     # via tensorflow
-tensorstore==0.1.28 \
-    --hash=sha256:197159da3ee97c08d8769fa151ff167328108bf07a13886f688b21f64abfb7d6 \
-    --hash=sha256:1c53e1e0e7996b8ba41854135d748c2bc8802b94a79122b56448aeb1f77efd9f \
-    --hash=sha256:6719eaba4ec2692d890c3c67c389fcae073d40d9e8c27e2dabd27fc9fd745e5c \
-    --hash=sha256:6bd122a44db81369808c84b7739ad2d8453f89c548c5d61d62f2b1f3331dbe9f \
-    --hash=sha256:7948ead0d5a55303ffa97a35659aa56a20c60bd89d9487643f51ac8d2e19b52c \
-    --hash=sha256:818aeb015b249069a686fb20f0c0b71663410ac6411cee88df18b8e05b567cbc \
-    --hash=sha256:b6f7e0c7455d9e164e6143714519e7c96cc4166b61ed7d268a8ab84a77e0d593 \
-    --hash=sha256:bd8615308ad3e29111238dfdd3aca1f3fb79b55c3dac1c04c96c8a319985addf \
-    --hash=sha256:c4de09e3fae30f8a9ab7648600c78e109ddd1d14af4280e055439fdd52a6dca9 \
-    --hash=sha256:cd8d8185136632c58edcd7cf4c43d301bf8ad61a197c632cb495d7d33b7be04e \
-    --hash=sha256:d7c86f63a6d4d7a84e3201ba9746ca51015612d375d77b9302eb93e9b88bc7aa \
-    --hash=sha256:dc82cbb77255fd2f7aa191f1e030f13276bf7272ea1027fcc6900537b30bfbcb \
-    --hash=sha256:deb32f1e74ab7b836ecec02a759a558fc57522a8dea0db4f19a01a78e93abab5 \
-    --hash=sha256:e480da4f95f7c95302c4b3e2a7fc32b8cfe1aa51742d72f84a335e49462a2465 \
-    --hash=sha256:ee16297a9493ed0d6d1000f216e79d29a0b80e7e97646e8520afd44138165839 \
-    --hash=sha256:f9270586401ee60ff79a4cd54c62ab5b06831d69b561df1fb9ebbeade4d4929c
-    # via flax
+tensorstore==0.1.41 \
+    --hash=sha256:025a62bb9122364885e90469af05fec2f62ad05f46ff46d9eae1d76ad9125563 \
+    --hash=sha256:2aa81581f768382a38584698a3fcb07a533fc391067467326656f24ab019cba1 \
+    --hash=sha256:2c4c578e82866b8f764de871ff7e0a81fe0949ac3565d8d2eb10f29e43020a52 \
+    --hash=sha256:2d65ea0fd5ac96a9d577f16bb917ae8a0a121d2093472bfb7bd762b1e32c753b \
+    --hash=sha256:5168f7f71e51da7d6cc85a11cd5d102d9eae750d5f5a3ee90cc9ebae10226621 \
+    --hash=sha256:620ad460023eeeae721e2e25a2a3e2b608f09cd169c1f68af7043c6d44e88cbf \
+    --hash=sha256:634c78fd62cd6e5357291ccb9671e43262f818f9cf7cc58f701b5bd80d1c1ef7 \
+    --hash=sha256:6b3b14616f9141b12e61c0c46d1c954927f7f307498d8b9d2261ff2bd4005bbd \
+    --hash=sha256:74e317ef7cba8c0208c5d8d9f1406eac37f58e8f92b3f7caa9a72b8b118b1c09 \
+    --hash=sha256:803fa2bcbc93f43fe0b3b5b70d78882d3b266a70d419acfc7fdd515f89cba79b \
+    --hash=sha256:8b5dbc0e809c90377527e0f65829d6abcdf5c69f892f433ed2cb8508d4ba519a \
+    --hash=sha256:8df13f990acc58889160eff5b2e1df029cdfffdf020ce5044e655242c1016bb1 \
+    --hash=sha256:91549a16b1ef2d6bc5ac8f28eed32737001fcfe33309f1ba126cd4c1e08b971b \
+    --hash=sha256:96fb62a880bf25da7e12ad4bba00a82deb2daf6f59050e8db6f0b04107120799 \
+    --hash=sha256:a99b87b65dfca65a830503bdfd2e5168a69b5290807cb8e922fa5a1acea2edec \
+    --hash=sha256:bbd58cedddce29216703a63ea42db010b6151c7bc05ac741af50aa31e31491fb \
+    --hash=sha256:c400aa46fc814edd69c72fcdf202dbd8c666ae684b534e81350a3a30ab16bdfc
+    # via
+    #   flax
+    #   orbax-checkpoint
 termcolor==2.1.1 \
     --hash=sha256:67cee2009adc6449c650f6bcf3bdeed00c8ba53a8cda5362733c53e0a39fb70b \
     --hash=sha256:fa852e957f97252205e105dd55bbc23b419a70fec0085708fc0515e399f304fd
@@ -1066,8 +865,8 @@ typing-extensions==4.4.0 \
     --hash=sha256:16fa4864408f655d35ec496218b85f79b3437c829e93320c7c9215ccfd92489e
     # via
     #   flax
-    #   jax
     #   optax
+    #   orbax-checkpoint
     #   rich
     #   tensorflow
 urllib3==1.26.13 \
diff --git a/tfjs-converter/python/requirements.txt b/tfjs-converter/python/requirements.txt
index 058dd2b6914..81b412a900f 100644
--- a/tfjs-converter/python/requirements.txt
+++ b/tfjs-converter/python/requirements.txt
@@ -1,8 +1,9 @@
-flax>=0.6.2,<0.6.3
-jax>=0.3.16
+flax>=0.7.2
 importlib_resources>=5.9.0
-tensorflow>=2.12.0,<3
-tensorflow-decision-forests>=1.3.0
-six>=1.12.0,<2
-tensorflow-hub>=0.13.0; python_version >= "3"
-packaging~=20.9
+jax>=0.4.13
+jaxlib>=0.4.13
+tensorflow>=2.13.0,<3
+tensorflow-decision-forests>=1.5.0
+six>=1.16.0,<2
+tensorflow-hub>=0.14.0
+packaging~=23.1
diff --git a/tfjs-converter/python/requirements_lock.txt b/tfjs-converter/python/requirements_lock.txt
index 081d58ba4c5..88cc454e85b 100644
--- a/tfjs-converter/python/requirements_lock.txt
+++ b/tfjs-converter/python/requirements_lock.txt
@@ -9,8 +9,8 @@ absl-py==1.3.0 \
     --hash=sha256:463c38a08d2e4cef6c498b76ba5bd4858e4c6ef51da1a5a1f27139a022e20248
     # via
     #   chex
-    #   jax
     #   optax
+    #   orbax-checkpoint
     #   tensorboard
     #   tensorflow
     #   tensorflow-decision-forests
@@ -18,6 +18,10 @@ astunparse==1.6.3 \
     --hash=sha256:5ad93a8456f0d084c3456d059fd9a92cce667963232cbf763eac3bc5b7940872 \
     --hash=sha256:c2652417f2c8b5bb325c885ae329bdf3f86424075c4fd1a128674bc6fba4b8e8
     # via tensorflow
+cached-property==1.5.2 \
+    --hash=sha256:9fa5755838eecbb2d234c3aa390bd80fbd3ac6b6869109bfc1b499f7bd89a130 \
+    --hash=sha256:df4f613cf7ad9a588cc381aaf4a512d26265ecebd5eb9e1ba12f1319eb85a6a0
+    # via orbax-checkpoint
 cachetools==5.2.0 \
     --hash=sha256:6a94c6402995a99c3970cc7e4884bb60b4a8639938157eeed436098bf9831757 \
     --hash=sha256:f9f17d2aec496a9aa6b76f53e3b614c965223c061982d434d160f930c698a9db
@@ -42,81 +46,6 @@ commonmark==0.9.1 \
     --hash=sha256:452f9dc859be7f06631ddcb328b6919c67984aca654e5fefb3914d54691aed60 \
     --hash=sha256:da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9
     # via rich
-contourpy==1.0.6 \
-    --hash=sha256:0236875c5a0784215b49d00ebbe80c5b6b5d5244b3655a36dda88105334dea17 \
-    --hash=sha256:03d1b9c6b44a9e30d554654c72be89af94fab7510b4b9f62356c64c81cec8b7d \
-    --hash=sha256:0537cc1195245bbe24f2913d1f9211b8f04eb203de9044630abd3664c6cc339c \
-    --hash=sha256:06ca79e1efbbe2df795822df2fa173d1a2b38b6e0f047a0ec7903fbca1d1847e \
-    --hash=sha256:08e8d09d96219ace6cb596506fb9b64ea5f270b2fb9121158b976d88871fcfd1 \
-    --hash=sha256:0b1e66346acfb17694d46175a0cea7d9036f12ed0c31dfe86f0f405eedde2bdd \
-    --hash=sha256:0b97454ed5b1368b66ed414c754cba15b9750ce69938fc6153679787402e4cdf \
-    --hash=sha256:0e4854cc02006ad6684ce092bdadab6f0912d131f91c2450ce6dbdea78ee3c0b \
-    --hash=sha256:12a7dc8439544ed05c6553bf026d5e8fa7fad48d63958a95d61698df0e00092b \
-    --hash=sha256:1b1ee48a130da4dd0eb8055bbab34abf3f6262957832fd575e0cab4979a15a41 \
-    --hash=sha256:1c0e1308307a75e07d1f1b5f0f56b5af84538a5e9027109a7bcf6cb47c434e72 \
-    --hash=sha256:1dedf4c64185a216c35eb488e6f433297c660321275734401760dafaeb0ad5c2 \
-    --hash=sha256:208bc904889c910d95aafcf7be9e677726df9ef71e216780170dbb7e37d118fa \
-    --hash=sha256:211dfe2bd43bf5791d23afbe23a7952e8ac8b67591d24be3638cabb648b3a6eb \
-    --hash=sha256:341330ed19074f956cb20877ad8d2ae50e458884bfa6a6df3ae28487cc76c768 \
-    --hash=sha256:344cb3badf6fc7316ad51835f56ac387bdf86c8e1b670904f18f437d70da4183 \
-    --hash=sha256:358f6364e4873f4d73360b35da30066f40387dd3c427a3e5432c6b28dd24a8fa \
-    --hash=sha256:371f6570a81dfdddbb837ba432293a63b4babb942a9eb7aaa699997adfb53278 \
-    --hash=sha256:375d81366afd547b8558c4720337218345148bc2fcffa3a9870cab82b29667f2 \
-    --hash=sha256:3a1917d3941dd58732c449c810fa7ce46cc305ce9325a11261d740118b85e6f3 \
-    --hash=sha256:4081918147fc4c29fad328d5066cfc751da100a1098398742f9f364be63803fc \
-    --hash=sha256:444fb776f58f4906d8d354eb6f6ce59d0a60f7b6a720da6c1ccb839db7c80eb9 \
-    --hash=sha256:46deb310a276cc5c1fd27958e358cce68b1e8a515fa5a574c670a504c3a3fe30 \
-    --hash=sha256:494efed2c761f0f37262815f9e3c4bb9917c5c69806abdee1d1cb6611a7174a0 \
-    --hash=sha256:50627bf76abb6ba291ad08db583161939c2c5fab38c38181b7833423ab9c7de3 \
-    --hash=sha256:5641927cc5ae66155d0c80195dc35726eae060e7defc18b7ab27600f39dd1fe7 \
-    --hash=sha256:5b117d29433fc8393b18a696d794961464e37afb34a6eeb8b2c37b5f4128a83e \
-    --hash=sha256:613c665529899b5d9fade7e5d1760111a0b011231277a0d36c49f0d3d6914bd6 \
-    --hash=sha256:6e459ebb8bb5ee4c22c19cc000174f8059981971a33ce11e17dddf6aca97a142 \
-    --hash=sha256:6f56515e7c6fae4529b731f6c117752247bef9cdad2b12fc5ddf8ca6a50965a5 \
-    --hash=sha256:730c27978a0003b47b359935478b7d63fd8386dbb2dcd36c1e8de88cbfc1e9de \
-    --hash=sha256:75a2e638042118118ab39d337da4c7908c1af74a8464cad59f19fbc5bbafec9b \
-    --hash=sha256:78ced51807ccb2f45d4ea73aca339756d75d021069604c2fccd05390dc3c28eb \
-    --hash=sha256:7ee394502026d68652c2824348a40bf50f31351a668977b51437131a90d777ea \
-    --hash=sha256:8468b40528fa1e15181cccec4198623b55dcd58306f8815a793803f51f6c474a \
-    --hash=sha256:84c593aeff7a0171f639da92cb86d24954bbb61f8a1b530f74eb750a14685832 \
-    --hash=sha256:913bac9d064cff033cf3719e855d4f1db9f1c179e0ecf3ba9fdef21c21c6a16a \
-    --hash=sha256:9447c45df407d3ecb717d837af3b70cfef432138530712263730783b3d016512 \
-    --hash=sha256:9b0e7fe7f949fb719b206548e5cde2518ffb29936afa4303d8a1c4db43dcb675 \
-    --hash=sha256:9bc407a6af672da20da74823443707e38ece8b93a04009dca25856c2d9adadb1 \
-    --hash=sha256:9e8e686a6db92a46111a1ee0ee6f7fbfae4048f0019de207149f43ac1812cf95 \
-    --hash=sha256:9fc4e7973ed0e1fe689435842a6e6b330eb7ccc696080dda9a97b1a1b78e41db \
-    --hash=sha256:a457ee72d9032e86730f62c5eeddf402e732fdf5ca8b13b41772aa8ae13a4563 \
-    --hash=sha256:a628bba09ba72e472bf7b31018b6281fd4cc903f0888049a3724afba13b6e0b8 \
-    --hash=sha256:a79d239fc22c3b8d9d3de492aa0c245533f4f4c7608e5749af866949c0f1b1b9 \
-    --hash=sha256:aa4674cf3fa2bd9c322982644967f01eed0c91bb890f624e0e0daf7a5c3383e9 \
-    --hash=sha256:acd2bd02f1a7adff3a1f33e431eb96ab6d7987b039d2946a9b39fe6fb16a1036 \
-    --hash=sha256:b3b1bd7577c530eaf9d2bc52d1a93fef50ac516a8b1062c3d1b9bcec9ebe329b \
-    --hash=sha256:b48d94386f1994db7c70c76b5808c12e23ed7a4ee13693c2fc5ab109d60243c0 \
-    --hash=sha256:b64f747e92af7da3b85631a55d68c45a2d728b4036b03cdaba4bd94bcc85bd6f \
-    --hash=sha256:b98c820608e2dca6442e786817f646d11057c09a23b68d2b3737e6dcb6e4a49b \
-    --hash=sha256:c1baa49ab9fedbf19d40d93163b7d3e735d9cd8d5efe4cce9907902a6dad391f \
-    --hash=sha256:c38c6536c2d71ca2f7e418acaf5bca30a3af7f2a2fa106083c7d738337848dbe \
-    --hash=sha256:c78bfbc1a7bff053baf7e508449d2765964d67735c909b583204e3240a2aca45 \
-    --hash=sha256:cd2bc0c8f2e8de7dd89a7f1c10b8844e291bca17d359373203ef2e6100819edd \
-    --hash=sha256:d2eff2af97ea0b61381828b1ad6cd249bbd41d280e53aea5cccd7b2b31b8225c \
-    --hash=sha256:d8834c14b8c3dd849005e06703469db9bf96ba2d66a3f88ecc539c9a8982e0ee \
-    --hash=sha256:d912f0154a20a80ea449daada904a7eb6941c83281a9fab95de50529bfc3a1da \
-    --hash=sha256:da1ef35fd79be2926ba80fbb36327463e3656c02526e9b5b4c2b366588b74d9a \
-    --hash=sha256:dbe6fe7a1166b1ddd7b6d887ea6fa8389d3f28b5ed3f73a8f40ece1fc5a3d340 \
-    --hash=sha256:dcd556c8fc37a342dd636d7eef150b1399f823a4462f8c968e11e1ebeabee769 \
-    --hash=sha256:e13b31d1b4b68db60b3b29f8e337908f328c7f05b9add4b1b5c74e0691180109 \
-    --hash=sha256:e1739496c2f0108013629aa095cc32a8c6363444361960c07493818d0dea2da4 \
-    --hash=sha256:e43255a83835a129ef98f75d13d643844d8c646b258bebd11e4a0975203e018f \
-    --hash=sha256:e626cefff8491bce356221c22af5a3ea528b0b41fbabc719c00ae233819ea0bf \
-    --hash=sha256:eadad75bf91897f922e0fb3dca1b322a58b1726a953f98c2e5f0606bd8408621 \
-    --hash=sha256:f33da6b5d19ad1bb5e7ad38bb8ba5c426d2178928bc2b2c44e8823ea0ecb6ff3 \
-    --hash=sha256:f4052a8a4926d4468416fc7d4b2a7b2a3e35f25b39f4061a7e2a3a2748c4fc48 \
-    --hash=sha256:f6ca38dd8d988eca8f07305125dec6f54ac1c518f1aaddcc14d08c01aebb6efc
-    # via matplotlib
-cycler==0.11.0 \
-    --hash=sha256:3a27e95f763a428a739d2add979fa7494c912a32c17c4c38c4d5f082cad165a3 \
-    --hash=sha256:9c87405839a19696e837b3b818fed3f5f69f16f1eec1a1ad77e043dcea9c772f
-    # via matplotlib
 dm-tree==0.1.7 \
     --hash=sha256:0f01743cc2247170e64798c6b4b31853717054bf9ceec47a1b1b8c2a4baf5792 \
     --hash=sha256:1379a02df36e2bbff9819ceafa55ccd436b15af398803f781f372f8ead7ed871 \
@@ -141,22 +70,18 @@ dm-tree==0.1.7 \
     --hash=sha256:d377bd621b485db42c4aeea0eabbd8f6274b89a9c338c2c1bf69a40c3b86a1fd \
     --hash=sha256:f3e2bd9b9c05d1a0039f7c128d8b055c8a05708ef569cdbbeec0a2946e425bd4
     # via chex
-etils[epath]==0.9.0 \
+etils==0.9.0 \
     --hash=sha256:489103e9e499a566765c60458ee15d185cf0065f2060a4d16a68f8f46962ed0d \
     --hash=sha256:635d6f7d1c519eb194304228543a4c5c7df0e6b58243302473e34c18cf720588
-    # via jax
-flatbuffers==22.12.6 \
-    --hash=sha256:27f67c6fb102d41c911c26867bda71f1f8622176ac072fa30f668f4d023b5826 \
-    --hash=sha256:b12a3214e73f325ecd392503e821ac43a791181356f020cc3657ba829b4947c8
+    # via orbax-checkpoint
+flatbuffers==23.5.26 \
+    --hash=sha256:9ea1144cac05ce5d86e2859f431c6cd5e66cd9c78c558317c7955fb8d4c78d89 \
+    --hash=sha256:c0ff356da363087b915fde4b8b45bdda73432fc17cddb3c8157472eab1422ad1
     # via tensorflow
-flax==0.6.2 \
-    --hash=sha256:9f933c87fb5762fbbf0920531e770bc385f24ef6eeb2f473641591fdbde9de89 \
-    --hash=sha256:a6247b412f14466fefcc70d043bd0facf72552d322acdda8a8700285308d390f
+flax==0.7.2 \
+    --hash=sha256:261c7b93e6d15ad80e2cedd2edb797d41b0b3c7805a54254de72a2366dc80148 \
+    --hash=sha256:7f023ece0b8b0d03019d2841dbe780e33b81ec42268bdc3e5d24c8fb0582fd7b
     # via -r tfjs-converter/python/requirements.txt
-fonttools==4.38.0 \
-    --hash=sha256:2bb244009f9bf3fa100fc3ead6aeb99febe5985fa20afbfbaa2f8946c2fbdaf1 \
-    --hash=sha256:820466f43c8be8c3009aef8b87e785014133508f0de64ec469e4efb643ae54fb
-    # via matplotlib
 gast==0.4.0 \
     --hash=sha256:40feb7b8b8434785585ab224d1568b857edb18297e5a3047f1ba012bc83b42c1 \
     --hash=sha256:b7adcdd5adbebf1adf17378da5ba3f543684dbec47b1cda1f3997e573cd542c4
@@ -254,112 +179,48 @@ idna==3.4 \
 importlib-metadata==5.1.0 \
     --hash=sha256:d5059f9f1e8e41f80e9c56c2ee58811450c31984dfa625329ffd7c0dad88a73b \
     --hash=sha256:d84d17e21670ec07990e1044a99efe8d615d860fd176fc29ef5c306068fda313
-    # via markdown
+    # via
+    #   jax
+    #   markdown
 importlib-resources==5.9.0 \
     --hash=sha256:5481e97fb45af8dcf2f798952625591c58fe599d0735d86b10f54de086a61681 \
     --hash=sha256:f78a8df21a79bcc30cfd400bdc38f314333de7c0fb619763f6b9dabab8268bb7
     # via
     #   -r tfjs-converter/python/requirements.txt
-    #   etils
-jax==0.3.16 \
-    --hash=sha256:211e6fb083fe846721285d5b9a6a22d16733f21a2ace350300039decfc7bf9a4
+    #   orbax-checkpoint
+jax==0.4.13 \
+    --hash=sha256:03bfe6749dfe647f16f15f6616638adae6c4a7ca7167c75c21961ecfd3a3baaa
     # via
     #   -r tfjs-converter/python/requirements.txt
     #   chex
     #   flax
     #   optax
-    #   tensorflow
-jaxlib==0.3.25 \
-    --hash=sha256:09508f7000c0fa958fba29267338e8de75b31d7ea29bd79719a568c38f0f8d31 \
-    --hash=sha256:13446a8382aa9ed944c16af636ca111d0afbbead91eed5cc2dc71195045e71b3 \
-    --hash=sha256:1e59ba58c9e93c1e1cef243f2609ec0b0c0a81160c20b9555aecdea32ccd6a78 \
-    --hash=sha256:1f1448f102a9d05186f579b6931fa0c607783ecc915fdfaa482c19538affa180 \
-    --hash=sha256:2e008e0a6c10aa7e949555e98dc0471e0d550d5d7c109771e38a971b49480538 \
-    --hash=sha256:3c75c8efd3702687968820446e3fb9ff997f8a2a07ab92e33b80e2f12eab3d9a \
-    --hash=sha256:5295354ed5db111e6f3e88cdfa4010d11c33dd926ac61735b9096b4e0746aa7b \
-    --hash=sha256:6e2f4e51041b8371aa3976b5a3a9cdcdccb1bd7b040c9b1345cbf24bd28a8d19 \
-    --hash=sha256:71866aeaafbc9a12b59dcbe443353772ef235b40c53f8bd7403d39311822c276 \
-    --hash=sha256:78b29c72d0680829db9377ed9be326875849258a60b8173b4a388b34ad18bc78 \
-    --hash=sha256:9f3116389ee834b3cdeb30001b085f4b55d7741366034f041c1d377154aa5afa \
-    --hash=sha256:a50193ba0cbf879021c9d73d7bcfa7eafb9138895d057b774c301aac3701f9a5 \
-    --hash=sha256:f2d517635fd77e2729c0ab7863be0d290927d01b2abb2f5dc955c821f8b0d53e \
-    --hash=sha256:fec043cdd55f3257d02e9d8880b33860eacadcae1bd5e26f43fdd08ada23614d
+    #   orbax-checkpoint
+jaxlib==0.4.13 \
+    --hash=sha256:19ae4c316b17a49342432c69f7f89f190b975333f3f9e9e175f686a651bc7347 \
+    --hash=sha256:411334d903df07dc1ace8d52fc53c17f6bc1d55aff7f6e0e5cf61ec149f758a0 \
+    --hash=sha256:49690fcdd26560515fd15399fc3a44777e0bfc5db5c48fe76ff7bc7228e8b2fb \
+    --hash=sha256:522635d5e159401a386c79f1236c218c1f68fbb4ca6648115c3ad3c2c3f518ab \
+    --hash=sha256:532ebc4fb11386282ad63b83941d4557f4038c1144acf026f1f8565f64c7e9c0 \
+    --hash=sha256:8000c0d15c107328e8f7b7b3ac91dd822f5c287a80231882b620503ed141fa89 \
+    --hash=sha256:839173b2e9593f5e9a6d3c42852cd15070fe80a939246efbb5cf40eec815de89 \
+    --hash=sha256:a259bb35429bfbd3b76e43019dfc8f7d6ea94bb217400b78f7d0824ce07a58ac \
+    --hash=sha256:b5c0a9737efd95fe18fd7715ce30dfce476546705ea8934aad6731777a9631a5 \
+    --hash=sha256:bebb4cf001f180dc431f9604daf930c2d9cc778e4dda26f401ac939b7bac912e \
+    --hash=sha256:c230ef85712e608d0f048869766a5a63afeb2e72309943db0df9f959ab17307f \
+    --hash=sha256:d19c05c15f962e098d49b45e2758aacf19330d192ec5395f9ef136f62db90edc \
+    --hash=sha256:ea1bc9811ef7d73a15e3213115e88fe7f5d14b59d95027bea9fccc98e5a14af8 \
+    --hash=sha256:f4e9e34e5d8a6556f62fead14aee0b1614c2c6296f0078d8e6139d6aff109649 \
+    --hash=sha256:fde66a93e9be89d99e5792f677ed8e319667d6b2396865b1c52c1312844c47f9
     # via
+    #   -r tfjs-converter/python/requirements.txt
     #   chex
     #   optax
-keras==2.12.0 \
-    --hash=sha256:35c39534011e909645fb93515452e98e1a0ce23727b55d4918b9c58b2308c15e
+    #   orbax-checkpoint
+keras==2.13.1 \
+    --hash=sha256:5ce5f706f779fa7330e63632f327b75ce38144a120376b2ae1917c00fa6136af \
+    --hash=sha256:5df12cc241a015a11b65ddb452c0eeb2744fce21d9b54ba48db87492568ccc68
     # via tensorflow
-kiwisolver==1.4.4 \
-    --hash=sha256:02f79693ec433cb4b5f51694e8477ae83b3205768a6fb48ffba60549080e295b \
-    --hash=sha256:03baab2d6b4a54ddbb43bba1a3a2d1627e82d205c5cf8f4c924dc49284b87166 \
-    --hash=sha256:1041feb4cda8708ce73bb4dcb9ce1ccf49d553bf87c3954bdfa46f0c3f77252c \
-    --hash=sha256:10ee06759482c78bdb864f4109886dff7b8a56529bc1609d4f1112b93fe6423c \
-    --hash=sha256:1d1573129aa0fd901076e2bfb4275a35f5b7aa60fbfb984499d661ec950320b0 \
-    --hash=sha256:283dffbf061a4ec60391d51e6155e372a1f7a4f5b15d59c8505339454f8989e4 \
-    --hash=sha256:28bc5b299f48150b5f822ce68624e445040595a4ac3d59251703779836eceff9 \
-    --hash=sha256:2a66fdfb34e05b705620dd567f5a03f239a088d5a3f321e7b6ac3239d22aa286 \
-    --hash=sha256:2e307eb9bd99801f82789b44bb45e9f541961831c7311521b13a6c85afc09767 \
-    --hash=sha256:2e407cb4bd5a13984a6c2c0fe1845e4e41e96f183e5e5cd4d77a857d9693494c \
-    --hash=sha256:2f5e60fabb7343a836360c4f0919b8cd0d6dbf08ad2ca6b9cf90bf0c76a3c4f6 \
-    --hash=sha256:36dafec3d6d6088d34e2de6b85f9d8e2324eb734162fba59d2ba9ed7a2043d5b \
-    --hash=sha256:3fe20f63c9ecee44560d0e7f116b3a747a5d7203376abeea292ab3152334d004 \
-    --hash=sha256:41dae968a94b1ef1897cb322b39360a0812661dba7c682aa45098eb8e193dbdf \
-    --hash=sha256:4bd472dbe5e136f96a4b18f295d159d7f26fd399136f5b17b08c4e5f498cd494 \
-    --hash=sha256:4ea39b0ccc4f5d803e3337dd46bcce60b702be4d86fd0b3d7531ef10fd99a1ac \
-    --hash=sha256:5853eb494c71e267912275e5586fe281444eb5e722de4e131cddf9d442615626 \
-    --hash=sha256:5bce61af018b0cb2055e0e72e7d65290d822d3feee430b7b8203d8a855e78766 \
-    --hash=sha256:6295ecd49304dcf3bfbfa45d9a081c96509e95f4b9d0eb7ee4ec0530c4a96514 \
-    --hash=sha256:62ac9cc684da4cf1778d07a89bf5f81b35834cb96ca523d3a7fb32509380cbf6 \
-    --hash=sha256:70e7c2e7b750585569564e2e5ca9845acfaa5da56ac46df68414f29fea97be9f \
-    --hash=sha256:7577c1987baa3adc4b3c62c33bd1118c3ef5c8ddef36f0f2c950ae0b199e100d \
-    --hash=sha256:75facbe9606748f43428fc91a43edb46c7ff68889b91fa31f53b58894503a191 \
-    --hash=sha256:787518a6789009c159453da4d6b683f468ef7a65bbde796bcea803ccf191058d \
-    --hash=sha256:78d6601aed50c74e0ef02f4204da1816147a6d3fbdc8b3872d263338a9052c51 \
-    --hash=sha256:7c43e1e1206cd421cd92e6b3280d4385d41d7166b3ed577ac20444b6995a445f \
-    --hash=sha256:81e38381b782cc7e1e46c4e14cd997ee6040768101aefc8fa3c24a4cc58e98f8 \
-    --hash=sha256:841293b17ad704d70c578f1f0013c890e219952169ce8a24ebc063eecf775454 \
-    --hash=sha256:872b8ca05c40d309ed13eb2e582cab0c5a05e81e987ab9c521bf05ad1d5cf5cb \
-    --hash=sha256:877272cf6b4b7e94c9614f9b10140e198d2186363728ed0f701c6eee1baec1da \
-    --hash=sha256:8c808594c88a025d4e322d5bb549282c93c8e1ba71b790f539567932722d7bd8 \
-    --hash=sha256:8ed58b8acf29798b036d347791141767ccf65eee7f26bde03a71c944449e53de \
-    --hash=sha256:91672bacaa030f92fc2f43b620d7b337fd9a5af28b0d6ed3f77afc43c4a64b5a \
-    --hash=sha256:968f44fdbf6dd757d12920d63b566eeb4d5b395fd2d00d29d7ef00a00582aac9 \
-    --hash=sha256:9f85003f5dfa867e86d53fac6f7e6f30c045673fa27b603c397753bebadc3008 \
-    --hash=sha256:a553dadda40fef6bfa1456dc4be49b113aa92c2a9a9e8711e955618cd69622e3 \
-    --hash=sha256:a68b62a02953b9841730db7797422f983935aeefceb1679f0fc85cbfbd311c32 \
-    --hash=sha256:abbe9fa13da955feb8202e215c4018f4bb57469b1b78c7a4c5c7b93001699938 \
-    --hash=sha256:ad881edc7ccb9d65b0224f4e4d05a1e85cf62d73aab798943df6d48ab0cd79a1 \
-    --hash=sha256:b1792d939ec70abe76f5054d3f36ed5656021dcad1322d1cc996d4e54165cef9 \
-    --hash=sha256:b428ef021242344340460fa4c9185d0b1f66fbdbfecc6c63eff4b7c29fad429d \
-    --hash=sha256:b533558eae785e33e8c148a8d9921692a9fe5aa516efbdff8606e7d87b9d5824 \
-    --hash=sha256:ba59c92039ec0a66103b1d5fe588fa546373587a7d68f5c96f743c3396afc04b \
-    --hash=sha256:bc8d3bd6c72b2dd9decf16ce70e20abcb3274ba01b4e1c96031e0c4067d1e7cd \
-    --hash=sha256:bc9db8a3efb3e403e4ecc6cd9489ea2bac94244f80c78e27c31dcc00d2790ac2 \
-    --hash=sha256:bf7d9fce9bcc4752ca4a1b80aabd38f6d19009ea5cbda0e0856983cf6d0023f5 \
-    --hash=sha256:c2dbb44c3f7e6c4d3487b31037b1bdbf424d97687c1747ce4ff2895795c9bf69 \
-    --hash=sha256:c79ebe8f3676a4c6630fd3f777f3cfecf9289666c84e775a67d1d358578dc2e3 \
-    --hash=sha256:c97528e64cb9ebeff9701e7938653a9951922f2a38bd847787d4a8e498cc83ae \
-    --hash=sha256:d0611a0a2a518464c05ddd5a3a1a0e856ccc10e67079bb17f265ad19ab3c7597 \
-    --hash=sha256:d06adcfa62a4431d404c31216f0f8ac97397d799cd53800e9d3efc2fbb3cf14e \
-    --hash=sha256:d41997519fcba4a1e46eb4a2fe31bc12f0ff957b2b81bac28db24744f333e955 \
-    --hash=sha256:d5b61785a9ce44e5a4b880272baa7cf6c8f48a5180c3e81c59553ba0cb0821ca \
-    --hash=sha256:da152d8cdcab0e56e4f45eb08b9aea6455845ec83172092f09b0e077ece2cf7a \
-    --hash=sha256:da7e547706e69e45d95e116e6939488d62174e033b763ab1496b4c29b76fabea \
-    --hash=sha256:db5283d90da4174865d520e7366801a93777201e91e79bacbac6e6927cbceede \
-    --hash=sha256:db608a6757adabb32f1cfe6066e39b3706d8c3aa69bbc353a5b61edad36a5cb4 \
-    --hash=sha256:e0ea21f66820452a3f5d1655f8704a60d66ba1191359b96541eaf457710a5fc6 \
-    --hash=sha256:e7da3fec7408813a7cebc9e4ec55afed2d0fd65c4754bc376bf03498d4e92686 \
-    --hash=sha256:e92a513161077b53447160b9bd8f522edfbed4bd9759e4c18ab05d7ef7e49408 \
-    --hash=sha256:ecb1fa0db7bf4cff9dac752abb19505a233c7f16684c5826d1f11ebd9472b871 \
-    --hash=sha256:efda5fc8cc1c61e4f639b8067d118e742b812c930f708e6667a5ce0d13499e29 \
-    --hash=sha256:f0a1dbdb5ecbef0d34eb77e56fcb3e95bbd7e50835d9782a45df81cc46949750 \
-    --hash=sha256:f0a71d85ecdd570ded8ac3d1c0f480842f49a40beb423bb8014539a9f32a5897 \
-    --hash=sha256:f4f270de01dd3e129a72efad823da90cc4d6aafb64c410c9033aba70db9f1ff0 \
-    --hash=sha256:f6cb459eea32a4e2cf18ba5fcece2dbdf496384413bc1bae15583f19e567f3b2 \
-    --hash=sha256:f8ad8285b01b0d4695102546b342b493b3ccc6781fc28c8c6a1bb63e95d22f09 \
-    --hash=sha256:f9f39e2f049db33a908319cf46624a569b36983c7c78318e9726a4cb8923b26c
-    # via matplotlib
 libclang==14.0.6 \
     --hash=sha256:206d2789e4450a37d054e63b70451a6fc1873466397443fa13de2b3d4adb2796 \
     --hash=sha256:2e4303e04517fcd11173cb2e51a7070eed71e16ef45d4e26a82c5e881cac3d27 \
@@ -418,49 +279,27 @@ markupsafe==2.1.1 \
     --hash=sha256:f121a1420d4e173a5d96e47e9a0c0dcff965afdf1626d28de1460815f7c4ee7a \
     --hash=sha256:fc7b548b17d238737688817ab67deebb30e8073c95749d55538ed473130ec0c7
     # via werkzeug
-matplotlib==3.6.2 \
-    --hash=sha256:0844523dfaaff566e39dbfa74e6f6dc42e92f7a365ce80929c5030b84caa563a \
-    --hash=sha256:0eda9d1b43f265da91fb9ae10d6922b5a986e2234470a524e6b18f14095b20d2 \
-    --hash=sha256:168093410b99f647ba61361b208f7b0d64dde1172b5b1796d765cd243cadb501 \
-    --hash=sha256:1836f366272b1557a613f8265db220eb8dd883202bbbabe01bad5a4eadfd0c95 \
-    --hash=sha256:19d61ee6414c44a04addbe33005ab1f87539d9f395e25afcbe9a3c50ce77c65c \
-    --hash=sha256:252957e208c23db72ca9918cb33e160c7833faebf295aaedb43f5b083832a267 \
-    --hash=sha256:32d29c8c26362169c80c5718ce367e8c64f4dd068a424e7110df1dd2ed7bd428 \
-    --hash=sha256:380d48c15ec41102a2b70858ab1dedfa33eb77b2c0982cb65a200ae67a48e9cb \
-    --hash=sha256:3964934731fd7a289a91d315919cf757f293969a4244941ab10513d2351b4e83 \
-    --hash=sha256:3cef89888a466228fc4e4b2954e740ce8e9afde7c4315fdd18caa1b8de58ca17 \
-    --hash=sha256:4426c74761790bff46e3d906c14c7aab727543293eed5a924300a952e1a3a3c1 \
-    --hash=sha256:5024b8ed83d7f8809982d095d8ab0b179bebc07616a9713f86d30cf4944acb73 \
-    --hash=sha256:52c2bdd7cd0bf9d5ccdf9c1816568fd4ccd51a4d82419cc5480f548981b47dd0 \
-    --hash=sha256:54fa9fe27f5466b86126ff38123261188bed568c1019e4716af01f97a12fe812 \
-    --hash=sha256:5ba73aa3aca35d2981e0b31230d58abb7b5d7ca104e543ae49709208d8ce706a \
-    --hash=sha256:5e16dcaecffd55b955aa5e2b8a804379789c15987e8ebd2f32f01398a81e975b \
-    --hash=sha256:5ecfc6559132116dedfc482d0ad9df8a89dc5909eebffd22f3deb684132d002f \
-    --hash=sha256:74153008bd24366cf099d1f1e83808d179d618c4e32edb0d489d526523a94d9f \
-    --hash=sha256:78ec3c3412cf277e6252764ee4acbdbec6920cc87ad65862272aaa0e24381eee \
-    --hash=sha256:795ad83940732b45d39b82571f87af0081c120feff2b12e748d96bb191169e33 \
-    --hash=sha256:7f716b6af94dc1b6b97c46401774472f0867e44595990fe80a8ba390f7a0a028 \
-    --hash=sha256:83dc89c5fd728fdb03b76f122f43b4dcee8c61f1489e232d9ad0f58020523e1c \
-    --hash=sha256:8a0ae37576ed444fe853709bdceb2be4c7df6f7acae17b8378765bd28e61b3ae \
-    --hash=sha256:8a8dbe2cb7f33ff54b16bb5c500673502a35f18ac1ed48625e997d40c922f9cc \
-    --hash=sha256:8a9d899953c722b9afd7e88dbefd8fb276c686c3116a43c577cfabf636180558 \
-    --hash=sha256:8d0068e40837c1d0df6e3abf1cdc9a34a6d2611d90e29610fa1d2455aeb4e2e5 \
-    --hash=sha256:9347cc6822f38db2b1d1ce992f375289670e595a2d1c15961aacbe0977407dfc \
-    --hash=sha256:9f335e5625feb90e323d7e3868ec337f7b9ad88b5d633f876e3b778813021dab \
-    --hash=sha256:b03fd10a1709d0101c054883b550f7c4c5e974f751e2680318759af005964990 \
-    --hash=sha256:b0ca2c60d3966dfd6608f5f8c49b8a0fcf76de6654f2eda55fc6ef038d5a6f27 \
-    --hash=sha256:b2604c6450f9dd2c42e223b1f5dca9643a23cfecc9fde4a94bb38e0d2693b136 \
-    --hash=sha256:ca0e7a658fbafcddcaefaa07ba8dae9384be2343468a8e011061791588d839fa \
-    --hash=sha256:d0e9ac04065a814d4cf2c6791a2ad563f739ae3ae830d716d54245c2b96fead6 \
-    --hash=sha256:d50e8c1e571ee39b5dfbc295c11ad65988879f68009dd281a6e1edbc2ff6c18c \
-    --hash=sha256:d840adcad7354be6f2ec28d0706528b0026e4c3934cc6566b84eac18633eab1b \
-    --hash=sha256:e0bbee6c2a5bf2a0017a9b5e397babb88f230e6f07c3cdff4a4c4bc75ed7c617 \
-    --hash=sha256:e5afe0a7ea0e3a7a257907060bee6724a6002b7eec55d0db16fd32409795f3e1 \
-    --hash=sha256:e68be81cd8c22b029924b6d0ee814c337c0e706b8d88495a617319e5dd5441c3 \
-    --hash=sha256:ec9be0f4826cdb3a3a517509dcc5f87f370251b76362051ab59e42b6b765f8c4 \
-    --hash=sha256:f04f97797df35e442ed09f529ad1235d1f1c0f30878e2fe09a2676b71a8801e0 \
-    --hash=sha256:f41e57ad63d336fe50d3a67bb8eaa26c09f6dda6a59f76777a99b8ccd8e26aec
-    # via flax
+ml-dtypes==0.2.0 \
+    --hash=sha256:022d5a4ee6be14569c2a9d1549e16f1ec87ca949681d0dca59995445d5fcdd5b \
+    --hash=sha256:1749b60348da71fd3c2ab303fdbc1965958dc50775ead41f5669c932a341cafd \
+    --hash=sha256:32107e7fa9f62db9a5281de923861325211dfff87bd23faefb27b303314635ab \
+    --hash=sha256:35b984cddbe8173b545a0e3334fe56ea1a5c3eb67c507f60d0cfde1d3fa8f8c2 \
+    --hash=sha256:36d28b8861a8931695e5a31176cad5ae85f6504906650dea5598fbec06c94606 \
+    --hash=sha256:50845af3e9a601810751b55091dee6c2562403fa1cb4e0123675cf3a4fc2c17a \
+    --hash=sha256:6488eb642acaaf08d8020f6de0a38acee7ac324c1e6e92ee0c0fea42422cb797 \
+    --hash=sha256:75015818a7fccf99a5e8ed18720cb430f3e71a8838388840f4cdf225c036c983 \
+    --hash=sha256:80d304c836d73f10605c58ccf7789c171cc229bfb678748adfb7cea2510dfd0e \
+    --hash=sha256:832a019a1b6db5c4422032ca9940a990fa104eee420f643713241b3a518977fa \
+    --hash=sha256:8faaf0897942c8253dd126662776ba45f0a5861968cf0f06d6d465f8a7bc298a \
+    --hash=sha256:bc29a0524ef5e23a7fbb8d881bdecabeb3fc1d19d9db61785d077a86cb94fab2 \
+    --hash=sha256:df6a76e1c8adf484feb138ed323f9f40a7b6c21788f120f7c78bec20ac37ee81 \
+    --hash=sha256:e70047ec2c83eaee01afdfdabee2c5b0c133804d90d0f7db4dd903360fcc537c \
+    --hash=sha256:e85ba8e24cf48d456e564688e981cf379d4c8e644db0a2f719b78de281bac2ca \
+    --hash=sha256:f00c71c8c63e03aff313bc6a7aeaac9a4f1483a921a6ffefa6d4404efd1af3d0 \
+    --hash=sha256:f08c391c2794f2aad358e6f4c70785a9a7b1df980ef4c232b3ccd4f6fe39f719
+    # via
+    #   jax
+    #   jaxlib
 msgpack==1.0.4 \
     --hash=sha256:002b5c72b6cd9b4bafd790f364b8480e859b4712e91f43014fe01e4f957b8467 \
     --hash=sha256:0a68d3ac0104e2d3510de90a1091720157c319ceeb90d74f7b5295a6bee51bae \
@@ -514,7 +353,13 @@ msgpack==1.0.4 \
     --hash=sha256:f5d869c18f030202eb412f08b28d2afeea553d6613aee89e200d7aca7ef01f5f \
     --hash=sha256:fb62ea4b62bfcb0b380d5680f9a4b3f9a2d166d9394e9bbd9666c0ee09a3645c \
     --hash=sha256:fcb8a47f43acc113e24e910399376f7277cf8508b27e5b88499f053de6b115a8
-    # via flax
+    # via
+    #   flax
+    #   orbax-checkpoint
+nest-asyncio==1.5.7 \
+    --hash=sha256:5301c82941b550b3123a1ea772ba9a1c80bad3a182be8c1a5ae6ad3be57a9657 \
+    --hash=sha256:6a80f7b98f24d9083ed24608977c09dd608d83f91cccc24c9d2cba6d10e01c10
+    # via orbax-checkpoint
 numpy==1.23.5 \
     --hash=sha256:01dd17cbb340bf0fc23981e52e1d18a9d4050792e8fb8363cecbf066a84b827d \
     --hash=sha256:06005a2ef6014e9956c09ba07654f9837d9e26696a0470e42beedadb78c11b07 \
@@ -546,14 +391,14 @@ numpy==1.23.5 \
     --hash=sha256:f9a909a8bae284d46bbfdefbdd4a262ba19d3bc9921b1e76126b1d21c3c34135
     # via
     #   chex
-    #   contourpy
     #   flax
     #   h5py
     #   jax
     #   jaxlib
-    #   matplotlib
+    #   ml-dtypes
     #   opt-einsum
     #   optax
+    #   orbax-checkpoint
     #   pandas
     #   scipy
     #   tensorboard
@@ -575,12 +420,15 @@ optax==0.1.4 \
     --hash=sha256:12fcf33bd682f9a162a3deb097f864130c3224d76771af2ba09410de80399a9b \
     --hash=sha256:fb7a0550d57a6636164a3de25986a8a19be8ff6431fcdf1225b4e05175810f22
     # via flax
-packaging==20.9 \
-    --hash=sha256:5b327ac1320dc863dca72f4514ecc086f31186744b84a230374cc1fd776feae5 \
-    --hash=sha256:67714da7f7bc052e064859c05c595155bd1ee9f69f76557e21f051443c20947a
+orbax-checkpoint==0.2.3 \
+    --hash=sha256:155e0a2dceef2901122e66585171e1dff4f4a4d9d2abe43a2b514279b9a3dabd \
+    --hash=sha256:a001bf48f1cebc635b07263fa546473ea48be3e278c50d5ade880b9aafb96f8a
+    # via flax
+packaging==23.1 \
+    --hash=sha256:994793af429502c4ea2ebf6bf664629d07c1a9fe974af92966e4b8d2df7edc61 \
+    --hash=sha256:a392980d2b6cffa644431898be54b0045151319d1e7ec34f0cfed48767dd334f
     # via
     #   -r tfjs-converter/python/requirements.txt
-    #   matplotlib
     #   tensorflow
 pandas==1.5.2 \
     --hash=sha256:0183cb04a057cc38fde5244909fca9826d5d57c4a5b7390c0cc3fa7acd9fa883 \
@@ -611,69 +459,6 @@ pandas==1.5.2 \
     --hash=sha256:e9dbacd22555c2d47f262ef96bb4e30880e5956169741400af8b306bbb24a273 \
     --hash=sha256:f6257b314fc14958f8122779e5a1557517b0f8e500cfb2bd53fa1f75a8ad0af2
     # via tensorflow-decision-forests
-pillow==9.3.0 \
-    --hash=sha256:03150abd92771742d4a8cd6f2fa6246d847dcd2e332a18d0c15cc75bf6703040 \
-    --hash=sha256:073adb2ae23431d3b9bcbcff3fe698b62ed47211d0716b067385538a1b0f28b8 \
-    --hash=sha256:0b07fffc13f474264c336298d1b4ce01d9c5a011415b79d4ee5527bb69ae6f65 \
-    --hash=sha256:0b7257127d646ff8676ec8a15520013a698d1fdc48bc2a79ba4e53df792526f2 \
-    --hash=sha256:12ce4932caf2ddf3e41d17fc9c02d67126935a44b86df6a206cf0d7161548627 \
-    --hash=sha256:15c42fb9dea42465dfd902fb0ecf584b8848ceb28b41ee2b58f866411be33f07 \
-    --hash=sha256:18498994b29e1cf86d505edcb7edbe814d133d2232d256db8c7a8ceb34d18cef \
-    --hash=sha256:1c7c8ae3864846fc95f4611c78129301e203aaa2af813b703c55d10cc1628535 \
-    --hash=sha256:22b012ea2d065fd163ca096f4e37e47cd8b59cf4b0fd47bfca6abb93df70b34c \
-    --hash=sha256:276a5ca930c913f714e372b2591a22c4bd3b81a418c0f6635ba832daec1cbcfc \
-    --hash=sha256:2e0918e03aa0c72ea56edbb00d4d664294815aa11291a11504a377ea018330d3 \
-    --hash=sha256:3033fbe1feb1b59394615a1cafaee85e49d01b51d54de0cbf6aa8e64182518a1 \
-    --hash=sha256:3168434d303babf495d4ba58fc22d6604f6e2afb97adc6a423e917dab828939c \
-    --hash=sha256:32a44128c4bdca7f31de5be641187367fe2a450ad83b833ef78910397db491aa \
-    --hash=sha256:3dd6caf940756101205dffc5367babf288a30043d35f80936f9bfb37f8355b32 \
-    --hash=sha256:40e1ce476a7804b0fb74bcfa80b0a2206ea6a882938eaba917f7a0f004b42502 \
-    --hash=sha256:41e0051336807468be450d52b8edd12ac60bebaa97fe10c8b660f116e50b30e4 \
-    --hash=sha256:4390e9ce199fc1951fcfa65795f239a8a4944117b5935a9317fb320e7767b40f \
-    --hash=sha256:502526a2cbfa431d9fc2a079bdd9061a2397b842bb6bc4239bb176da00993812 \
-    --hash=sha256:51e0e543a33ed92db9f5ef69a0356e0b1a7a6b6a71b80df99f1d181ae5875636 \
-    --hash=sha256:57751894f6618fd4308ed8e0c36c333e2f5469744c34729a27532b3db106ee20 \
-    --hash=sha256:5d77adcd56a42d00cc1be30843d3426aa4e660cab4a61021dc84467123f7a00c \
-    --hash=sha256:655a83b0058ba47c7c52e4e2df5ecf484c1b0b0349805896dd350cbc416bdd91 \
-    --hash=sha256:68943d632f1f9e3dce98908e873b3a090f6cba1cbb1b892a9e8d97c938871fbe \
-    --hash=sha256:6c738585d7a9961d8c2821a1eb3dcb978d14e238be3d70f0a706f7fa9316946b \
-    --hash=sha256:73bd195e43f3fadecfc50c682f5055ec32ee2c933243cafbfdec69ab1aa87cad \
-    --hash=sha256:772a91fc0e03eaf922c63badeca75e91baa80fe2f5f87bdaed4280662aad25c9 \
-    --hash=sha256:77ec3e7be99629898c9a6d24a09de089fa5356ee408cdffffe62d67bb75fdd72 \
-    --hash=sha256:7db8b751ad307d7cf238f02101e8e36a128a6cb199326e867d1398067381bff4 \
-    --hash=sha256:801ec82e4188e935c7f5e22e006d01611d6b41661bba9fe45b60e7ac1a8f84de \
-    --hash=sha256:82409ffe29d70fd733ff3c1025a602abb3e67405d41b9403b00b01debc4c9a29 \
-    --hash=sha256:828989c45c245518065a110434246c44a56a8b2b2f6347d1409c787e6e4651ee \
-    --hash=sha256:829f97c8e258593b9daa80638aee3789b7df9da5cf1336035016d76f03b8860c \
-    --hash=sha256:871b72c3643e516db4ecf20efe735deb27fe30ca17800e661d769faab45a18d7 \
-    --hash=sha256:89dca0ce00a2b49024df6325925555d406b14aa3efc2f752dbb5940c52c56b11 \
-    --hash=sha256:90fb88843d3902fe7c9586d439d1e8c05258f41da473952aa8b328d8b907498c \
-    --hash=sha256:97aabc5c50312afa5e0a2b07c17d4ac5e865b250986f8afe2b02d772567a380c \
-    --hash=sha256:9aaa107275d8527e9d6e7670b64aabaaa36e5b6bd71a1015ddd21da0d4e06448 \
-    --hash=sha256:9f47eabcd2ded7698106b05c2c338672d16a6f2a485e74481f524e2a23c2794b \
-    --hash=sha256:a0a06a052c5f37b4ed81c613a455a81f9a3a69429b4fd7bb913c3fa98abefc20 \
-    --hash=sha256:ab388aaa3f6ce52ac1cb8e122c4bd46657c15905904b3120a6248b5b8b0bc228 \
-    --hash=sha256:ad58d27a5b0262c0c19b47d54c5802db9b34d38bbf886665b626aff83c74bacd \
-    --hash=sha256:ae5331c23ce118c53b172fa64a4c037eb83c9165aba3a7ba9ddd3ec9fa64a699 \
-    --hash=sha256:af0372acb5d3598f36ec0914deed2a63f6bcdb7b606da04dc19a88d31bf0c05b \
-    --hash=sha256:afa4107d1b306cdf8953edde0534562607fe8811b6c4d9a486298ad31de733b2 \
-    --hash=sha256:b03ae6f1a1878233ac620c98f3459f79fd77c7e3c2b20d460284e1fb370557d4 \
-    --hash=sha256:b0915e734b33a474d76c28e07292f196cdf2a590a0d25bcc06e64e545f2d146c \
-    --hash=sha256:b4012d06c846dc2b80651b120e2cdd787b013deb39c09f407727ba90015c684f \
-    --hash=sha256:b472b5ea442148d1c3e2209f20f1e0bb0eb556538690fa70b5e1f79fa0ba8dc2 \
-    --hash=sha256:b59430236b8e58840a0dfb4099a0e8717ffb779c952426a69ae435ca1f57210c \
-    --hash=sha256:b90f7616ea170e92820775ed47e136208e04c967271c9ef615b6fbd08d9af0e3 \
-    --hash=sha256:b9a65733d103311331875c1dca05cb4606997fd33d6acfed695b1232ba1df193 \
-    --hash=sha256:bac18ab8d2d1e6b4ce25e3424f709aceef668347db8637c2296bcf41acb7cf48 \
-    --hash=sha256:bca31dd6014cb8b0b2db1e46081b0ca7d936f856da3b39744aef499db5d84d02 \
-    --hash=sha256:be55f8457cd1eac957af0c3f5ece7bc3f033f89b114ef30f710882717670b2a8 \
-    --hash=sha256:c7025dce65566eb6e89f56c9509d4f628fddcedb131d9465cacd3d8bac337e7e \
-    --hash=sha256:c935a22a557a560108d780f9a0fc426dd7459940dc54faa49d83249c8d3e760f \
-    --hash=sha256:dbb8e7f2abee51cef77673be97760abff1674ed32847ce04b4af90f610144c7b \
-    --hash=sha256:e6ea6b856a74d560d9326c0f5895ef8050126acfdc7ca08ad703eb0081e82b74 \
-    --hash=sha256:ebf2029c1f464c59b8bdbe5143c79fa2045a581ac53679733d3a91d400ff9efb \
-    --hash=sha256:f1ff2ee69f10f13a9596480335f406dd1f70c3650349e2be67ca3139280cade0
-    # via matplotlib
 protobuf==4.22.3 \
     --hash=sha256:13233ee2b9d3bd9a5f216c1fa2c321cd564b93d8f2e4f521a85b585447747997 \
     --hash=sha256:23452f2fdea754a8251d0fc88c0317735ae47217e0d27bf330a30eec2848811a \
@@ -706,18 +491,10 @@ pygments==2.13.0 \
     --hash=sha256:56a8508ae95f98e2b9bdf93a6be5ae3f7d8af858b43e02c5a2ff083726be40c1 \
     --hash=sha256:f643f331ab57ba3c9d89212ee4a2dabc6e94f117cf4eefde99a0574720d14c42
     # via rich
-pyparsing==3.0.9 \
-    --hash=sha256:2b020ecf7d21b687f219b71ecad3631f644a47f01403fa1d1036b0c6416d70fb \
-    --hash=sha256:5026bae9a10eeaefb61dab2f09052b9f4307d44aee4eda64b309723d8d206bbc
-    # via
-    #   matplotlib
-    #   packaging
 python-dateutil==2.8.2 \
     --hash=sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86 \
     --hash=sha256:961d03dc3453ebbc59dbdea9e4e11c5651520a876d0f4db161e8674aae935da9
-    # via
-    #   matplotlib
-    #   pandas
+    # via pandas
 pytz==2022.6 \
     --hash=sha256:222439474e9c98fced559f1709d89e6c9cbf8d79c794ff3eb9f8800064291427 \
     --hash=sha256:e89512406b793ca39f5971bc999cc538ce125c0e51c27941bef4568b460095e2
@@ -763,7 +540,9 @@ pyyaml==6.0 \
     --hash=sha256:dbad0e9d368bb989f4515da330b88a057617d16b6a8245084f1b05400f24609f \
     --hash=sha256:e61ceaab6f49fb8bdfaa0f92c4b57bcfbea54c09277b1b4f7ac376bfb7a7c174 \
     --hash=sha256:f84fbc98b019fef2ee9a1cb3ce93e3187a6df0b2538a651bfb890254ba9f90b5
-    # via flax
+    # via
+    #   flax
+    #   orbax-checkpoint
 requests==2.28.1 \
     --hash=sha256:7c5599b102feddaa661c826c56ab4fee28bfd17f5abca1ebbe3e7f19d7c97983 \
     --hash=sha256:8fefa2a1a1365bf5520aac41836fbee479da67864514bdb821f31ce07ce65349
@@ -813,9 +592,9 @@ setuptools==65.6.3 \
     # via
     #   tensorboard
     #   tensorflow
-six==1.12.0 \
-    --hash=sha256:3350809f0555b11f552448330d0b52d5f24c91a322ea4a15ef22629740f3761c \
-    --hash=sha256:d16a0141ec1a18405cd4ce8b4613101da75da0e9a7aec5bdd4fa804d0e0eba73
+six==1.16.0 \
+    --hash=sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926 \
+    --hash=sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254
     # via
     #   -r tfjs-converter/python/requirements.txt
     #   astunparse
@@ -824,47 +603,57 @@ six==1.12.0 \
     #   python-dateutil
     #   tensorflow
     #   tensorflow-decision-forests
-tensorboard==2.12.2 \
-    --hash=sha256:811ab0d27a139445836db9fd4f974424602c3dce12379364d379bcba7c783a68
+tensorboard==2.13.0 \
+    --hash=sha256:ab69961ebddbddc83f5fa2ff9233572bdad5b883778c35e4fe94bf1798bd8481
     # via tensorflow
 tensorboard-data-server==0.7.0 \
     --hash=sha256:64aa1be7c23e80b1a42c13b686eb0875bb70f5e755f4d2b8de5c1d880cf2267f \
     --hash=sha256:753d4214799b31da7b6d93837959abebbc6afa86e69eacf1e9a317a48daa31eb \
     --hash=sha256:eb7fa518737944dbf4f0cf83c2e40a7ac346bf91be2e6a0215de98be74e85454
     # via tensorboard
-tensorboard-plugin-wit==1.8.1 \
-    --hash=sha256:ff26bdd583d155aa951ee3b152b3d0cffae8005dc697f72b44a8e8c2a77a8cbe
-    # via tensorboard
-tensorflow==2.12.0 \
-    --hash=sha256:020d6a54cb26020bdc71a7bae8ee35be05096f63e773dc517f6e87c49de62c50 \
-    --hash=sha256:23850332f1f9f778d697c9dba63ca52be72cb73363e75ad358f07ddafef63c01 \
-    --hash=sha256:31f81eb8adaeb558963f5d8b47dbfcc398d898f0857bf3de6b6484350236b7b5 \
-    --hash=sha256:357d9d2851188a8d27ee195345b4d175cad970150d1344ba9d9fcc4bf2b68336 \
-    --hash=sha256:42fc2635e9420faee781a16bd393126f29cd39aa2b9d02901f24d8497bd6f958 \
-    --hash=sha256:4afc2dd57435f29ebe249eb5f595d89b0e73be94922eeb7110aa6280a332837c \
-    --hash=sha256:6e7641e2a6e32f31ff233495478a9cc86b7c038140eab714a61eeddbbbb327c3 \
-    --hash=sha256:6ec4a2934ea19e92f27a9668ece43025ed5efe14b5d19be53b07692bc8a4189d \
-    --hash=sha256:76414355e420edb9154b4e72113eef5813ccb71701fda959afbbc1eebe3099bd \
-    --hash=sha256:91dccda42c03569d8c787190482a11ecae3b9b173aaa9166f0ab20cecc9c31f4 \
-    --hash=sha256:9f70a8f9ab46e5ed436850aa60d1cd40645f5c669e14bcad48915dc1f597dda2 \
-    --hash=sha256:a7194e744c5a7f3e759ecb949527b4a07718a6d1110e6e82fd4ce0c5586a7d4a \
-    --hash=sha256:be4ac0dfcc7a16f6df2bc19bd322e312235ab3f7b0c7297f96c92c44bb14d2a1 \
-    --hash=sha256:c5193ddb3bb5120cb445279beb08ed9e74a85a4eeb2485550d6fb707a89d9a88 \
-    --hash=sha256:c8001210df7202ef6267150865b0b79f834c3ca69ee3132277de8eeb994dffde \
-    --hash=sha256:e29fcf6cfd069aefb4b44f357cccbb4415a5a3d7b5b516eaf4450062fe40021e
+tensorflow==2.13.0 \
+    --hash=sha256:00060c5516a61e30c51936084ebc37091d116efe9ae74b2818cbd8b2006218e7 \
+    --hash=sha256:06559eeaa69e6561cccbe2d02b015bcec663e875c8bbc4643f55692476e52147 \
+    --hash=sha256:076d953a1508dc58bf95f30f58bcc9ee364b1353c61e143cb20c2dada91afb05 \
+    --hash=sha256:11ad6a7ff49b4a690fbf37a5adaf28ba2686350a859c5f13c58dc8d2cc670375 \
+    --hash=sha256:19ee67901702b26787ad685cca65730c163c101c0c2f238a2584d714e0fa8c25 \
+    --hash=sha256:2822ac48c38d69b7fb104e606dacbd763c4bf5d3b20791f25be16a5076d01912 \
+    --hash=sha256:5e0fdadec59de3d11c5b5129ddc38e739bde7aa13095b82e19d4380e14d04999 \
+    --hash=sha256:6fff426661d286a4c634da44275d2ea2b951f392f3e65c8603681e7cb040586a \
+    --hash=sha256:72d68b8c2f382e2d01b956c8ba516c0a7d5dad98111dd351bf82bfa646aa1c72 \
+    --hash=sha256:7a08c0e2938ed5b642a8787678123827477b81d316055d5073fff82fa183eb82 \
+    --hash=sha256:89125443e998548059c4e4a129dfab2b1ff7f2fd4c8eaed97842c3cd9b663101 \
+    --hash=sha256:948003b5a23b72b3d89746d729e62ec5f01e47460f05521b2211d95069f569ba \
+    --hash=sha256:9c04bc3023b6c4cfb9ee9759c3f03f21993891b4c345df52eb5519204fbf28c0 \
+    --hash=sha256:b2978b39e8b3919059b5fd9e28508d50a77965d06ed0b537ed71c97de22dabdf \
+    --hash=sha256:cbb83561bb7d55859eaefc70c674e58713d4e10c10927423ed836a5289bbfa86 \
+    --hash=sha256:de77306c0c22c9d8754f54700752ac3a1efee895c5357308e6594436404bfbc0 \
+    --hash=sha256:e0cf94d36ceaba8f158c6e15404a81fd5b3aa4cb04147c674cf55bd1aec78154 \
+    --hash=sha256:e8f0b69ee2f800399fc6bc7ec55fecfa33662d136e425485959d90638f32a32a \
+    --hash=sha256:fa7abe265cc3ebccc9b405a280bf674824c6d85df5e6ccfa985987b3c9d265b4 \
+    --hash=sha256:fb2ff1129c93e853c19897d6a22ed0ec56387f5c6290ec03dec1c6f7b80bc396
     # via
     #   -r tfjs-converter/python/requirements.txt
     #   tensorflow-decision-forests
-tensorflow-decision-forests==1.3.0 \
-    --hash=sha256:9729973a6b16ee75161cf2e8d68a8fddef3988db8e05734b94a3fad6780a916d \
-    --hash=sha256:bb8b07d3e39e00153428288202e1393991308411c0af542afdc7b721856aa15a \
-    --hash=sha256:dbd9b8563af9ad5d56381fe483d75e593c8ecb820f1f671a79cc74fafc906225
+tensorflow-decision-forests==1.5.0 \
+    --hash=sha256:04fd913627d08fe54514b179c612e87eebf55f1448bf01951660985dfa14a6e1 \
+    --hash=sha256:1209a2832ac65f8f74bd9d0c1d58f3f8b771e7fa5c9d504c547842311647b7d4 \
+    --hash=sha256:22e3835acbfbd5356bb2f8e0c973dfc40ef80a7924b793e90c811158448cfe77 \
+    --hash=sha256:43ffd4fba1c3376f58a9dcee943df80f0cff6e47224d109ad0389a723c74947c \
+    --hash=sha256:4a0df3a3be5751594d49f5a8f99977b553cf1c42f320b952ac2a2f67b85283f5 \
+    --hash=sha256:5ec4297eb5e7c4110cf8aae89e9b08b9ad2cb725e3e63c89c78304c0d7235d24 \
+    --hash=sha256:804f6bed277b5c5b6d2bd85738a64973d5d3e8e6ac06abf6098545740245cedc \
+    --hash=sha256:a43af2a5a8c34e550bf549c6cad96da271979efc5a8ec988f6f76cc90770415a \
+    --hash=sha256:d137241dad8e884d0c937aa8769fe0768324804e9ba666a78b7b5f2f536a0bd2 \
+    --hash=sha256:d685e92abe44920ee6d89394ec4e075bb1ada7402f673566146e1b476a576e96 \
+    --hash=sha256:f5d8c3730578bda55a8f520ae39b0c9b2560d69bd53b57882e5371c1a82ba098 \
+    --hash=sha256:fbd403acf736bb9b4afd2985d9056e6d5043fc4b9a31bd05e5fcae2b1d413dc3
     # via -r tfjs-converter/python/requirements.txt
-tensorflow-estimator==2.12.0 \
-    --hash=sha256:59b191bead4883822de3d63ac02ace11a83bfe6c10d64d0c4dfde75a50e60ca1
+tensorflow-estimator==2.13.0 \
+    --hash=sha256:6f868284eaa654ae3aa7cacdbef2175d0909df9fcf11374f5166f8bf475952aa
     # via tensorflow
-tensorflow-hub==0.13.0 ; python_version >= "3" \
-    --hash=sha256:3544f4fd9fd99e4eeb6da1b5b5320e4a2dbdef7f9bb778f66f76d6790f32dd65
+tensorflow-hub==0.14.0 \
+    --hash=sha256:519c6b56c4d304667fbd8ce66bd637e6a750c901215468db2cc6bfd0739bb0b0
     # via -r tfjs-converter/python/requirements.txt
 tensorflow-io-gcs-filesystem==0.28.0 \
     --hash=sha256:00cf6a92f1f9f90b2ba2d728870bcd2a70b116316d0817ab0b91dd390c25b3fd \
@@ -883,24 +672,27 @@ tensorflow-io-gcs-filesystem==0.28.0 \
     --hash=sha256:cc062ce13ec95fb64b1fd426818a6d2b0e5be9692bc0e43a19cce115b6da4336 \
     --hash=sha256:f76cbe1a784841c223f6861e5f6c7e53aa6232cb626d57e76881a0638c365de6
     # via tensorflow
-tensorstore==0.1.28 \
-    --hash=sha256:197159da3ee97c08d8769fa151ff167328108bf07a13886f688b21f64abfb7d6 \
-    --hash=sha256:1c53e1e0e7996b8ba41854135d748c2bc8802b94a79122b56448aeb1f77efd9f \
-    --hash=sha256:6719eaba4ec2692d890c3c67c389fcae073d40d9e8c27e2dabd27fc9fd745e5c \
-    --hash=sha256:6bd122a44db81369808c84b7739ad2d8453f89c548c5d61d62f2b1f3331dbe9f \
-    --hash=sha256:7948ead0d5a55303ffa97a35659aa56a20c60bd89d9487643f51ac8d2e19b52c \
-    --hash=sha256:818aeb015b249069a686fb20f0c0b71663410ac6411cee88df18b8e05b567cbc \
-    --hash=sha256:b6f7e0c7455d9e164e6143714519e7c96cc4166b61ed7d268a8ab84a77e0d593 \
-    --hash=sha256:bd8615308ad3e29111238dfdd3aca1f3fb79b55c3dac1c04c96c8a319985addf \
-    --hash=sha256:c4de09e3fae30f8a9ab7648600c78e109ddd1d14af4280e055439fdd52a6dca9 \
-    --hash=sha256:cd8d8185136632c58edcd7cf4c43d301bf8ad61a197c632cb495d7d33b7be04e \
-    --hash=sha256:d7c86f63a6d4d7a84e3201ba9746ca51015612d375d77b9302eb93e9b88bc7aa \
-    --hash=sha256:dc82cbb77255fd2f7aa191f1e030f13276bf7272ea1027fcc6900537b30bfbcb \
-    --hash=sha256:deb32f1e74ab7b836ecec02a759a558fc57522a8dea0db4f19a01a78e93abab5 \
-    --hash=sha256:e480da4f95f7c95302c4b3e2a7fc32b8cfe1aa51742d72f84a335e49462a2465 \
-    --hash=sha256:ee16297a9493ed0d6d1000f216e79d29a0b80e7e97646e8520afd44138165839 \
-    --hash=sha256:f9270586401ee60ff79a4cd54c62ab5b06831d69b561df1fb9ebbeade4d4929c
-    # via flax
+tensorstore==0.1.41 \
+    --hash=sha256:025a62bb9122364885e90469af05fec2f62ad05f46ff46d9eae1d76ad9125563 \
+    --hash=sha256:2aa81581f768382a38584698a3fcb07a533fc391067467326656f24ab019cba1 \
+    --hash=sha256:2c4c578e82866b8f764de871ff7e0a81fe0949ac3565d8d2eb10f29e43020a52 \
+    --hash=sha256:2d65ea0fd5ac96a9d577f16bb917ae8a0a121d2093472bfb7bd762b1e32c753b \
+    --hash=sha256:5168f7f71e51da7d6cc85a11cd5d102d9eae750d5f5a3ee90cc9ebae10226621 \
+    --hash=sha256:620ad460023eeeae721e2e25a2a3e2b608f09cd169c1f68af7043c6d44e88cbf \
+    --hash=sha256:634c78fd62cd6e5357291ccb9671e43262f818f9cf7cc58f701b5bd80d1c1ef7 \
+    --hash=sha256:6b3b14616f9141b12e61c0c46d1c954927f7f307498d8b9d2261ff2bd4005bbd \
+    --hash=sha256:74e317ef7cba8c0208c5d8d9f1406eac37f58e8f92b3f7caa9a72b8b118b1c09 \
+    --hash=sha256:803fa2bcbc93f43fe0b3b5b70d78882d3b266a70d419acfc7fdd515f89cba79b \
+    --hash=sha256:8b5dbc0e809c90377527e0f65829d6abcdf5c69f892f433ed2cb8508d4ba519a \
+    --hash=sha256:8df13f990acc58889160eff5b2e1df029cdfffdf020ce5044e655242c1016bb1 \
+    --hash=sha256:91549a16b1ef2d6bc5ac8f28eed32737001fcfe33309f1ba126cd4c1e08b971b \
+    --hash=sha256:96fb62a880bf25da7e12ad4bba00a82deb2daf6f59050e8db6f0b04107120799 \
+    --hash=sha256:a99b87b65dfca65a830503bdfd2e5168a69b5290807cb8e922fa5a1acea2edec \
+    --hash=sha256:bbd58cedddce29216703a63ea42db010b6151c7bc05ac741af50aa31e31491fb \
+    --hash=sha256:c400aa46fc814edd69c72fcdf202dbd8c666ae684b534e81350a3a30ab16bdfc
+    # via
+    #   flax
+    #   orbax-checkpoint
 termcolor==2.1.1 \
     --hash=sha256:67cee2009adc6449c650f6bcf3bdeed00c8ba53a8cda5362733c53e0a39fb70b \
     --hash=sha256:fa852e957f97252205e105dd55bbc23b419a70fec0085708fc0515e399f304fd
@@ -913,10 +705,9 @@ typing-extensions==4.4.0 \
     --hash=sha256:1511434bb92bf8dd198c12b1cc812e800d4181cfcb867674e0f8279cc93087aa \
     --hash=sha256:16fa4864408f655d35ec496218b85f79b3437c829e93320c7c9215ccfd92489e
     # via
-    #   etils
     #   flax
-    #   jax
     #   optax
+    #   orbax-checkpoint
     #   tensorflow
 urllib3==1.26.13 \
     --hash=sha256:47cc05d99aaa09c9e72ed5809b60e7ba354e64b59c9c173ac3018642d8bb41fc \
@@ -1007,6 +798,5 @@ zipp==3.11.0 \
     --hash=sha256:83a28fcb75844b5c0cdaf5aa4003c2d728c77e05f5aeabe8e95e56727005fbaa \
     --hash=sha256:a7a22e05929290a67401440b39690ae6563279bced5f314609d9d03798f56766
     # via
-    #   etils
     #   importlib-metadata
     #   importlib-resources
diff --git a/tfjs-converter/python/tensorflowjs/converters/converter_test.py b/tfjs-converter/python/tensorflowjs/converters/converter_test.py
index 25b95a19dcd..c1b237e456c 100644
--- a/tfjs-converter/python/tensorflowjs/converters/converter_test.py
+++ b/tfjs-converter/python/tensorflowjs/converters/converter_test.py
@@ -100,7 +100,6 @@ def testConvertSavedKerasModelNoSplitByLayer(self):
     self.assertIn('layers', model_json['model_config']['config'])
 
     # Check the loaded weights.
-    self.assertEqual(tf.keras.__version__, model_json['keras_version'])
     self.assertEqual('tensorflow', model_json['backend'])
     self.assertEqual(1, len(groups))
     self.assertEqual(3, len(groups[0]))
@@ -136,7 +135,6 @@ def testConvertSavedKerasModelSplitByLayer(self):
     self.assertIn('layers', model_json['model_config']['config'])
 
     # Check the loaded weights.
-    self.assertEqual(tf.keras.__version__, model_json['keras_version'])
     self.assertEqual('tensorflow', model_json['backend'])
     self.assertEqual(2, len(groups))
     self.assertEqual(2, len(groups[0]))
@@ -274,6 +272,7 @@ def testTensorflowjsToKerasConversionSucceeds(self):
           tf.keras.layers.Dense(
               1, use_bias=False, kernel_initializer='ones', name='Dense2')])
       h5_path = os.path.join(self._tmp_dir, 'SequentialModel.h5')
+      sequential_model.use_legacy_config = True
       sequential_model.save(h5_path)
       converter.dispatch_keras_h5_to_tfjs_layers_model_conversion(
           h5_path, output_dir=self._tmp_dir)
@@ -287,6 +286,7 @@ def testTensorflowjsToKerasConversionSucceeds(self):
     # Load the new H5 and compare the model JSONs.
     with tf.Graph().as_default(), tf.compat.v1.Session():
       new_model = tf.keras.models.load_model(new_h5_path)
+      new_model.use_legacy_config = True
       self.assertEqual(old_model_json, new_model.to_json())
 
   def testTensorflowjsToKerasConversionFailsOnDirInputPath(self):
diff --git a/tfjs-converter/python/tensorflowjs/converters/keras_h5_conversion_test.py b/tfjs-converter/python/tensorflowjs/converters/keras_h5_conversion_test.py
index 5732fccbfb1..e4e0b6c1417 100644
--- a/tfjs-converter/python/tensorflowjs/converters/keras_h5_conversion_test.py
+++ b/tfjs-converter/python/tensorflowjs/converters/keras_h5_conversion_test.py
@@ -135,8 +135,6 @@ def testConvertModelWithNestedLayerNames(self):
 
     # Check meta-data in the artifact JSON.
     self.assertEqual(model_json['format'], 'layers-model')
-    self.assertEqual(model_json['generatedBy'],
-                     'keras v%s' % tf.keras.__version__)
     self.assertEqual(
         model_json['convertedBy'],
         'TensorFlow.js Converter v%s' % version.version)
@@ -191,7 +189,6 @@ def testConvertMergedModelFromSimpleModelNoSplitByLayer(self):
     # By default, all weights of the model ought to be put in the same group.
     self.assertEqual(1, len(groups))
 
-    self.assertEqual(tf.keras.__version__, out['keras_version'])
     self.assertEqual('tensorflow', out['backend'])
     weight_group = groups[0]
     self.assertEqual(3, len(weight_group))
@@ -244,7 +241,6 @@ def testConvertMergedModelFromSimpleModelSplitByLayer(self):
     # because the model has two layers.
     self.assertEqual(2, len(groups))
 
-    self.assertEqual(tf.keras.__version__, out['keras_version'])
     self.assertEqual('tensorflow', out['backend'])
     self.assertEqual(2, len(groups[0]))
     kernel1 = groups[0][0]
diff --git a/tfjs-converter/python/tensorflowjs/converters/keras_tfjs_loader_test.py b/tfjs-converter/python/tensorflowjs/converters/keras_tfjs_loader_test.py
index 21cf3fd6e97..a9a166d716c 100755
--- a/tfjs-converter/python/tensorflowjs/converters/keras_tfjs_loader_test.py
+++ b/tfjs-converter/python/tensorflowjs/converters/keras_tfjs_loader_test.py
@@ -46,6 +46,7 @@ def tearDown(self):
 
   def _saveKerasModelForTest(self, path):
     model = tf.keras.Sequential()
+    model.use_legacy_config = True
     model.add(tf.keras.layers.Dense(
         2, input_shape=[12], bias_initializer='random_normal', name='dense'))
     model.add(tf.keras.layers.Dense(
@@ -58,6 +59,7 @@ def _saveKerasModelForTest(self, path):
 
   def _saveRNNKerasModelForTest(self, path):
     model = tf.keras.Sequential()
+    model.use_legacy_config = True
     model.add(tf.keras.layers.Embedding(100, 20, input_shape=[10]))
     model.add(tf.keras.layers.SimpleRNN(4))
     keras_h5_conversion.save_keras_model(model, path)
@@ -75,7 +77,7 @@ def testLoadKerasModelAndWeights(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.load_keras_model(
           os.path.join(tfjs_path, 'model.json'))
-
+      model2.use_legacy_config = True
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))
@@ -98,7 +100,7 @@ def testLoadKerasRNNModelAndWeights(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.load_keras_model(
           os.path.join(tfjs_path, 'model.json'))
-
+      model2.use_legacy_config = True
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))
@@ -127,6 +129,7 @@ def testDeserializeKerasModelTopologyOnlyFromBytesIO(self):
 
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.deserialize_keras_model(buff.read())
+      model2.use_legacy_config = True
 
       # The two model JSONs should match exactly.
       self.assertEqual(model1.to_json(), model2.to_json())
@@ -145,7 +148,7 @@ def testDeserializeKerasModelTopologyOnlyFromJSONDict(self):
 
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.deserialize_keras_model(config_json)
-
+      model2.use_legacy_config = True
       # The two model JSONs should match exactly.
       self.assertEqual(model1.to_json(), model2.to_json())
 
@@ -171,6 +174,7 @@ def testDeserializeKerasModelTopologyAndWeightsFromBuffers(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.deserialize_keras_model(
           json_buff, weight_data=weight_buffers)
+      model2.use_legacy_config = True
 
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
@@ -200,6 +204,7 @@ def testDeserializeKerasModelTopologyAndWeightsFromFileObjects(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.deserialize_keras_model(
           json_file, weight_files)
+      model2.use_legacy_config = True
 
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
@@ -225,7 +230,7 @@ def testLoadKerasModelWithCurrentWorkingDirectoryRelativePath(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       # Use a relative path under the current working directory.
       model2 = keras_tfjs_loader.load_keras_model('model.json')
-
+      model2.use_legacy_config = True
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))
@@ -247,6 +252,7 @@ def testLoadKerasModelWithoutWeights(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.load_keras_model(
           os.path.join(tfjs_path, 'model.json'), load_weights=False)
+      model2.use_legacy_config = True
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))
       for model1_weight_value, model2_weight_value in zip(
@@ -276,7 +282,7 @@ def testLoadKerasModelFromNonDefaultWeightsPathWorks(self):
     with tf.Graph().as_default(), tf.compat.v1.Session():
       model2 = keras_tfjs_loader.load_keras_model(
           new_model_json_path, weights_path_prefix=tfjs_path)
-
+      model2.use_legacy_config = True
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))
@@ -329,7 +335,7 @@ def testLoadKerasModelFromDataBuffers(self):
       model2 = keras_tfjs_loader.load_keras_model(
           os.path.join(tfjs_path, 'model.json'),
           weights_data_buffers=data_buffers)
-
+      model2.use_legacy_config = True
       # Verify the equality of all the weight values.
       model2_weight_values = model2.get_weights()
       self.assertEqual(len(model1_weight_values), len(model2_weight_values))

From d00c13d72c54e6920afd1abbb56fad6aed197a94 Mon Sep 17 00:00:00 2001
From: Fengwu Yao <fengwuyao@google.com>
Date: Fri, 11 Aug 2023 12:54:47 -0700
Subject: [PATCH 2/3] fix jax test

---
 .../python/tensorflowjs/converters/jax_conversion_test.py       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tfjs-converter/python/tensorflowjs/converters/jax_conversion_test.py b/tfjs-converter/python/tensorflowjs/converters/jax_conversion_test.py
index 8f08e8e08e0..cf7e49a14c3 100644
--- a/tfjs-converter/python/tensorflowjs/converters/jax_conversion_test.py
+++ b/tfjs-converter/python/tensorflowjs/converters/jax_conversion_test.py
@@ -97,7 +97,7 @@ def test_convert_poly(self):
   def test_convert_tf_poly_mismatch_raises(self):
     apply_fn = lambda params, x: jnp.sum(x) * params['w']
     with self.assertRaisesRegex(
-        ValueError, 'polymorphic shape.* must match .* for argument shape'):
+        ValueError, 'syntax error in polymorphic shape.* in dimension.* Parsed.*, remaining.*'):
       jax_conversion.convert_jax(
           apply_fn,
           {'w': 0.5},

From d0b43349918d9310c9d88ca60e2f725827421202 Mon Sep 17 00:00:00 2001
From: Fengwu Yao <fengwuyao@google.com>
Date: Fri, 11 Aug 2023 15:40:43 -0700
Subject: [PATCH 3/3] fix version check tests

---
 .../python/tensorflowjs/converters/converter_test.py       | 4 +++-
 .../tensorflowjs/converters/keras_h5_conversion_test.py    | 7 +++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/tfjs-converter/python/tensorflowjs/converters/converter_test.py b/tfjs-converter/python/tensorflowjs/converters/converter_test.py
index c1b237e456c..397e94d2ced 100644
--- a/tfjs-converter/python/tensorflowjs/converters/converter_test.py
+++ b/tfjs-converter/python/tensorflowjs/converters/converter_test.py
@@ -24,6 +24,7 @@
 import shutil
 import tempfile
 import unittest
+import keras
 
 import numpy as np
 import tensorflow.compat.v2 as tf
@@ -98,8 +99,8 @@ def testConvertSavedKerasModelNoSplitByLayer(self):
     self.assertIsInstance(model_json['model_config'], dict)
     self.assertIsInstance(model_json['model_config']['config'], dict)
     self.assertIn('layers', model_json['model_config']['config'])
-
     # Check the loaded weights.
+    self.assertEqual(keras.__version__, model_json['keras_version'])
     self.assertEqual('tensorflow', model_json['backend'])
     self.assertEqual(1, len(groups))
     self.assertEqual(3, len(groups[0]))
@@ -135,6 +136,7 @@ def testConvertSavedKerasModelSplitByLayer(self):
     self.assertIn('layers', model_json['model_config']['config'])
 
     # Check the loaded weights.
+    self.assertEqual(keras.__version__, model_json['keras_version'])
     self.assertEqual('tensorflow', model_json['backend'])
     self.assertEqual(2, len(groups))
     self.assertEqual(2, len(groups[0]))
diff --git a/tfjs-converter/python/tensorflowjs/converters/keras_h5_conversion_test.py b/tfjs-converter/python/tensorflowjs/converters/keras_h5_conversion_test.py
index e4e0b6c1417..f3c95e3aad8 100644
--- a/tfjs-converter/python/tensorflowjs/converters/keras_h5_conversion_test.py
+++ b/tfjs-converter/python/tensorflowjs/converters/keras_h5_conversion_test.py
@@ -25,6 +25,7 @@
 import tempfile
 import unittest
 import six
+import keras
 
 import h5py
 import numpy as np
@@ -135,6 +136,8 @@ def testConvertModelWithNestedLayerNames(self):
 
     # Check meta-data in the artifact JSON.
     self.assertEqual(model_json['format'], 'layers-model')
+    self.assertEqual(model_json['generatedBy'],
+                     'keras v%s' % keras.__version__)
     self.assertEqual(
         model_json['convertedBy'],
         'TensorFlow.js Converter v%s' % version.version)
@@ -188,7 +191,7 @@ def testConvertMergedModelFromSimpleModelNoSplitByLayer(self):
     # Check the loaded weights.
     # By default, all weights of the model ought to be put in the same group.
     self.assertEqual(1, len(groups))
-
+    self.assertEqual(keras.__version__, out['keras_version'])
     self.assertEqual('tensorflow', out['backend'])
     weight_group = groups[0]
     self.assertEqual(3, len(weight_group))
@@ -240,7 +243,7 @@ def testConvertMergedModelFromSimpleModelSplitByLayer(self):
     # Due to `split_by_layer=True`, there ought to be two weight groups,
     # because the model has two layers.
     self.assertEqual(2, len(groups))
-
+    self.assertEqual(keras.__version__, out['keras_version'])
     self.assertEqual('tensorflow', out['backend'])
     self.assertEqual(2, len(groups[0]))
     kernel1 = groups[0][0]