docstr: Runs a training loop, asserting that callbacks are called.
docstr_tok: ["Runs","a","training","loop",",","asserting","that","callbacks","are","called","."]
code: def test_calls_step_callbacks(self):
      """Runs a training loop, asserting that callbacks are called."""
      call_at_steps = [1, 3, 4]
      begin_steps = []
      end_steps = []
      test_case = self
  
      class TestCallback(callbacks.TrainingStepCallback):
          def call_at(self, step):
              return step in call_at_steps
  
          def on_step_begin(self, step):
              begin_steps.append(step)
  
          def on_step_end(self, step):
              # Assert that on_step_begin() was called before.
              test_case.assertIn(step, begin_steps)
              end_steps.append(step)
  
      model = tl.Serial(tl.Dense(1))
      task = training.TrainTask(
          _very_simple_data(), tl.L2Loss(), optimizers.SGD(0.01)
      )
      loop = training.Loop(model, [task], callbacks=[TestCallback])
      loop.run(n_steps=5)
  
      # Assert that the callback has been called at the appropriate steps.
      self.assertEqual(begin_steps, call_at_steps)
      self.assertEqual(end_steps, call_at_steps)
  
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 26 for calling the TrainingTest obj's underlying member method code with example usage: obj.test_calls_step_callbacks() and returns: unknown