class_docstr: Test cases for train.
class_docstr_tok: ["Test","cases","for","train","."]
docstr: Tests CNN module used as the trainable model.
docstr_tok: ["Tests","CNN","module","used","as","the","trainable","model","."]
code: def test_cnn(self):
      """Tests CNN module used as the trainable model."""
      rng = jax.random.key(0)
      inputs = jnp.ones((1, 28, 28, 3), jnp.float32)
      output, variables = train.CNN().init_with_output(rng, inputs)
  
      self.assertEqual((1, 10), output.shape)
      self.assertEqual(
          CNN_PARAMS,
          sum(
              np.prod(arr.shape)
              for arr in jax.tree_util.tree_leaves(variables["params"])
          ),
      )
  
file_path: flax/examples/mnist/train_test.py
imports: import pathlib
  import tempfile
  from absl.testing import absltest
  import jax
  from jax import numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
  from configs import default
  import train
label_desc: Use image node_id 2 for calling the TrainTest obj's underlying member method code with example usage: obj.test_cnn() without return types