docstr: Saves and restores a checkpoint to check for equivalence.
docstr_tok: ["Saves","and","restores","a","checkpoint","to","check","for","equivalence","."]
code: def test_train_save_restore_dense(self):
      """Saves and restores a checkpoint to check for equivalence."""
      self.skipTest(
          "Broken by https://github.com/google/jax/pull/11234"
      )
      train_data = data.Serial(
          lambda _: _very_simple_data(),
          data.CountAndSkip("simple_data"),
      )
      task = training.TrainTask(
          train_data(), tl.L2Loss(), optimizers.Adam(0.0001)
      )
      eval_task = training.EvalTask(
          _very_simple_data(),  # deliberately re-using training data
          [tl.L2Loss()],
          metric_names=["SGD.L2Loss"],
      )
      tmp_dir = self.create_tempdir().full_path
  
      def _make_model_and_session():
          m = tl.Serial(tl.Dense(1))
          ts = training.Loop(
              m,
              [task],
              eval_tasks=[eval_task],
              eval_at=lambda step_n: step_n % 2 == 0,
              output_dir=tmp_dir,
          )
          return m, ts
  
      model, training_session = _make_model_and_session()
      self.assertEqual(0, training_session.step)
      training_session.run(n_steps=1)
      training_session.save_checkpoint("model")
      self.assertEqual(data.inputs.data_counters["simple_data"], 2)
      data.inputs.data_counters["simple_data"] = 0  # reset manually
      self.assertEqual(
          data.inputs.data_counters["simple_data"], 0
      )  # check
      model2, training_session2 = _make_model_and_session()
      self.assertEqual(
          data.inputs.data_counters["simple_data"], 2
      )  # restored
  
      x = np.ones((8, 1))
      y1 = model(x, rng=fastmath.random.get_prng(0))
      y2 = model2(x, rng=fastmath.random.get_prng(0))
      self.assertEqual(str(y1), str(y2))
  
      training_session2.run(n_steps=1)
      y1 = model(x, rng=fastmath.random.get_prng(0))
      y2 = model2(x, rng=fastmath.random.get_prng(0))
      self.assertNotEqual(str(y1), str(y2))
  
      slots1 = training_session._trainer_per_task[0].slots
      slots2 = training_session2._trainer_per_task[0].slots
      np.testing.assert_array_equal(slots1, slots2)
  
imports: import collections
  import os
  import time
  from absl.testing import absltest
  from jax import config
  import numpy
  from trax import data
  from trax import fastmath
  from trax import layers
  from trax import optimizers
  from trax import shapes
  from trax import test_utils
  from trax.layers import base
  from trax.models import transformer
  from trax.supervised import callbacks
  from trax.supervised import training
label_desc: Use image node_id 7 for calling the TrainingTest obj's underlying member method code with example usage: obj.test_train_save_restore_dense() and returns: m,  ts