def test_restores_step_bfloat16(self):
    """Training restores step from directory where it saved it, w/ bfloat16."""
    self.skipTest('Broken by https://github.com/google/jax/pull/11234')
    model = tl.Serial(tl.Dense(1, use_bfloat16=True))
    # We'll also use Adafactor with bfloat16 to check restoring bfloat slots.
    opt = optimizers.Adafactor(.01, do_momentum=True, momentum_in_bfloat16=True)
    task = training.TrainTask(_very_simple_data(), tl.L2Loss(), opt)
    tmp_dir = self.create_tempdir().full_path
    loop = training.Loop(model, [task],
                         checkpoint_at=lambda step_n: step_n % 2 == 0,
                         output_dir=tmp_dir)
    loop.run(4)
    loop2 = training.Loop(model, [task], output_dir=tmp_dir)
    self.assertEqual(4, loop2.step)
    loop2.run(2)  # check that continued training works
    self.assertEqual(6, loop2.step)
Training restores step from directory where it saved it, w/ bfloat16.
trax/supervised/training_test.py