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(.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)
Saves and restores a checkpoint to check for equivalence.
trax/supervised/training_test.py