JAX on the Web with TensorFlow.js

<!-- Loads global symbol `lit`. --> <script src="exports_bin.js"></script> <script> async function demo() {   lit.setBaseUrl('https://google-research.github.io/vision_transformer/lit');   const model = new lit.Model('tiny');   await model.load();   console.log(model.computeProbabilities(['a dog', 'a cat'], /*imgIdx=*/1); } demo(); </script>

Conclusion

In this article you learned how to convert JAX functions and Flax models into the TensorFlow.js format that can be executed in a browser or on devices capable of running JavaScript.

The first example demonstrated how to convert a JAX function to a TensorFlow.js model, which can then be loaded in Colab for verification, or run on any device with a modern web browser – this is an exactly the same conversion that can be applied to more complex Flax models. The second example showed how to train an ML model in Colab, and test it interactively on a mobile phone.The third example provided a ). We hope that this application can serve you as a good starting point for your own client-side demos using JAX models with TensorFlow.js.