class_docstr: An MNIST model using Flax.
docstr: Trains a pure JAX MNIST predictor.
  
  Returns:
    a tuple with two elements:
      - a predictor function with signature "(Params, ImagesBatch) ->
        Predictions".
        If `with_classifier=False` then the output of the predictor function
        is the last layer of logits.
      - the parameters "Params" for the predictor function
docstr_tok: ["Trains","a","pure","JAX","MNIST","predictor",".","Returns",":","a","tuple","with","two","elements",":","-","a","predictor","function","with","signature","``","(","Params",",","ImagesBatch",")","-",">","Predictions","''",".","If","`","with_classifier=False","`","then","the","output","of","the","predictor","function","is","the","last","layer","of","logits",".","-","the","parameters","``","Params","''","for","the","predictor","function"]
code: def train(train_ds, test_ds, num_epochs, with_classifier=True):
      """Trains a pure JAX MNIST predictor.
  
      Returns:
        a tuple with two elements:
          - a predictor function with signature "(Params, ImagesBatch) ->
            Predictions".
            If `with_classifier=False` then the output of the predictor function
            is the last layer of logits.
          - the parameters "Params" for the predictor function
      """
      rng = jax.random.PRNGKey(0)
      momentum_mass = 0.9
  
      init_shape = jnp.ones((1,) + input_shape, jnp.float32)
      params = FlaxMNIST.model.init(rng, init_shape)["params"]
      tx = optax.sgd(learning_rate=step_size, momentum=momentum_mass)
      opt_state = tx.init(params)
  
      for epoch in range(num_epochs):
          start_time = time.time()
          for inputs, labels in tfds.as_numpy(train_ds):
              params, opt_state = jax.jit(
                  FlaxMNIST.update, static_argnums=0
              )(tx, params, opt_state, inputs, labels)
          epoch_time = time.time() - start_time
          # Same accuracy function as for the pure JAX example
          train_acc = PureJaxMNIST.accuracy(
              FlaxMNIST.predict, params, train_ds
          )
          test_acc = PureJaxMNIST.accuracy(
              FlaxMNIST.predict, params, test_ds
          )
          logging.info(
              "%s: Epoch %d in %0.2f sec",
              FlaxMNIST.name,
              epoch,
              epoch_time,
          )
          logging.info(
              "%s: Training set accuracy %0.2f%%",
              FlaxMNIST.name,
              100.0 * train_acc,
          )
          logging.info(
              "%s: Test set accuracy %0.2f%%",
              FlaxMNIST.name,
              100.0 * test_acc,
          )
  
      # See discussion in README.md for packaging Flax models for conversion
      predict_fn = functools.partial(
          FlaxMNIST.predict, with_classifier=with_classifier
      )
      return (predict_fn, params)
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  import logging
  import re
  import time
  from typing import Any, Callable, Optional
  from absl import flags
  import flax
  from flax import linen
  import jax
  import jax.numpy
  from matplotlib import pyplot
  import numpy
  import optax
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 4 for calling the FlaxMNIST obj's underlying member method code with example usage: obj.train(train_ds, test_ds, num_epochs, with_classifier) and returns: predict_fn,  params