JAX Linear Model

In this example, we show how to optimize JAX linear model. We recommend running this example in NVIDIA NGC TensorFlow2 container.

Set up the environment

. ./install.sh

Run optimization

./optimize.py