Temporarily pin flax to <0.6.3 to fix circular dep build issue

[mattsoulanille]

Flax 0.6.3 seems to cause a circular dependency issue with Bazel. Pin it to <0.6.3 for now.

ERROR: /usr/local/google/home/msoulanille/.cache/bazel/_bazel_msoulanille/9633cfa0c28af5f24d900dd13e6bc5b8/external/tensorflowjs_dev_deps/pypi__flax/BUILD.bazel:22:11: in py_library rule @tensorflowjs_dev_deps//pypi__flax:pypi__flax: cycle in dependency graph:
    //tfjs-converter/python/tensorflowjs/converters:jax_conversion_test (1ac704f01a0ba260bfb1a01d72a5e91827144d78836864b5710509d3d9d956a2)
    //tfjs-converter/python/tensorflowjs:expect_flax_installed (c57eef692cab50e751ca54fdd64a9d8b7f624d9bccb12be15778f5c2d8865e6c)
.-> @tensorflowjs_dev_deps//pypi__flax:pypi__flax (c57eef692cab50e751ca54fdd64a9d8b7f624d9bccb12be15778f5c2d8865e6c)
|   @tensorflowjs_dev_deps//pypi__orbax:pypi__orbax (c57eef692cab50e751ca54fdd64a9d8b7f624d9bccb12be15778f5c2d8865e6c)
`-- @tensorflowjs_dev_deps//pypi__flax:pypi__flax (c57eef692cab50e751ca54fdd64a9d8b7f624d9bccb12be15778f5c2d8865e6c)

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

---

This change is 

[mattsoulanille]

Superseded by #7160