def test_eval_equals_predict_configs(inp, model_fn, configs, seq_axis=1,
                                     seq_tensor=None, message=''):
  """Utility method for testing equivalence of predict and eval modes.

  This function iterates over a list of dictionaries `confis`, and runs the test
  on models with each configuration.

  Args:
    inp: input fed to the model. It can be a tensor, or a tuple of tensors.
    model_fn: function creating a model after calling with `mode` argument.
    configs: List of dictionaries, which contain configs to be fed into
      `model_fn`.
    seq_axis: axis of sequence_length. In predict mode we iterate over this
      axis. By default `1`, which is 2nd dimension.
    seq_tensor: if `inp` is a tuple, `seq_tensor` is an index of an input tensor
      in this tuple on which we iterate the sequence.
    message: Optional message to show when outputs of eval/predict mode don't
      match.
  """
  for config in configs:
    model_fn_configured = functools.partial(model_fn, **config)
    test_eval_equals_predict(inp, model_fn_configured, seq_axis=seq_axis,
                             seq_tensor=seq_tensor,
                             message=' Config: {}.{}'.format(config, message))
Utility method for testing equivalence of predict and eval modes.

  This function iterates over a list of dictionaries `confis`, and runs the test
  on models with each configuration.

  Args:
    inp: input fed to the model. It can be a tensor, or a tuple of tensors.
    model_fn: function creating a model after calling with `mode` argument.
    configs: List of dictionaries, which contain configs to be fed into
      `model_fn`.
    seq_axis: axis of sequence_length. In predict mode we iterate over this
      axis. By default `1`, which is 2nd dimension.
    seq_tensor: if `inp` is a tuple, `seq_tensor` is an index of an input tensor
      in this tuple on which we iterate the sequence.
    message: Optional message to show when outputs of eval/predict mode don't
      match.
  
trax/layers/test_utils.py