def AssertFunction(specification, layer, message=None):  # pylint: disable=invalid-name
  """AssertFunction asserts shapes on the input/output tensors of a layer.

  It passes all inputs to the layer, and returns all outputs of the layer
  unchanged.

  Args:
    specification: A specification. See assert_shape decorator for a full
        documentation.
    layer: A base.Layer to wrap around.
    message: An optional message to print if an assert fails. By default it will
        print the filename and the line number where AssertFunction was called.

  Returns:
    The given layer wrapped in asserts on its inputs and outputs.
  """
  if message is None:
    caller = inspect.getframeinfo(inspect.stack()[1][0])
    message = f'Defined at {caller.filename}:{caller.lineno}'
  before_spec, after_spec = specification.split('->')
  before_assert = AssertShape(before_spec, message=message + ' function input')
  after_assert = AssertShape(after_spec, message=message + ' function output')
  after_assert._create_link(before_assert)  # pylint: disable=protected-access
  return combinators.Serial(
      before_assert, layer, after_assert)
AssertFunction asserts shapes on the input/output tensors of a layer.

  It passes all inputs to the layer, and returns all outputs of the layer
  unchanged.

  Args:
    specification: A specification. See assert_shape decorator for a full
        documentation.
    layer: A base.Layer to wrap around.
    message: An optional message to print if an assert fails. By default it will
        print the filename and the line number where AssertFunction was called.

  Returns:
    The given layer wrapped in asserts on its inputs and outputs.
  
trax/layers/assert_shape.py