code: 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)
  
file_path: trax/trax/layers/assert_shape.py
imports: import functools
  import inspect
  import string
  from absl import logging
  from trax.layers import base
  from trax.layers import combinators
label_desc: Use image node_id 2 for calling a global function with example usage: AssertFunction(specification, layer, message) and returns: combinators