This model's TensorFlow and Flax versions
were contributed by ydshieh.