class_docstr: Benchmarks for the MNIST Flax example.
class_docstr_tok: ["Benchmarks","for","the","MNIST","Flax","example","."]
docstr: Run full training for MNIST CPU training.
docstr_tok: ["Run","full","training","for","MNIST","CPU","training","."]
code: def test_cpu(self):
      """Run full training for MNIST CPU training."""
      # Prepare and set flags defined in main.py.
      workdir = self.get_tmp_model_dir()
      config = default.get_config()
  
      FLAGS.workdir = workdir
      FLAGS.config = config
  
      start_time = time.time()
      main.main([])
      benchmark_time = time.time() - start_time
  
      summaries = self.read_summaries(workdir)
  
      # Summaries contain all the information necessary for the regression
      # metrics.
      wall_time, _, eval_accuracy = zip(*summaries["eval_accuracy"])
      wall_time = np.array(wall_time)
      sec_per_epoch = np.mean(wall_time[1:] - wall_time[:-1])
      end_eval_accuracy = eval_accuracy[-1]
  
      # Assertions are deferred until the test finishes, so the metrics are
      # always reported and benchmark success is determined based on *all*
      # assertions.
      self.assertBetween(end_eval_accuracy, 0.98, 1.0)
  
      # Use the reporting API to report single or multiple metrics/extras.
      self.report_wall_time(benchmark_time)
      self.report_metrics(
          {
              "sec_per_epoch": sec_per_epoch,
              "accuracy": end_eval_accuracy,
          }
      )
      self.report_extras(
          {
              "model_name": "MNIST",
              "description": "CPU test for MNIST.",
              "implementation": "linen",
          }
      )
  
file_path: flax/examples/mnist/mnist_benchmark.py
imports: import time
  from absl import flags
  from absl.testing import absltest
  from absl.testing.flagsaver import flagsaver
  from flax.testing import Benchmark
  import jax
  import numpy
  import main
  from configs import default
label_desc: Use image node_id 1 for calling the MnistBenchmark obj's underlying member method code with example usage: obj.test_cpu() without return types