code: def assert_shape(specification):
      """Decorator for checking the input and output shapes of Layer.
  
      Decorator can be applied on trax.base.Layer class, or a function returning
      a trax.base.Layer class. It uses notation similar to einsum (Einstein
      summation convention), achieving concise and simple representation of tensors.
      For example 'ij,jh->ih' is a valid representation of a function taking two
      2D matrices as input, and returning a single output, also a 2D matrix.
  
      It improves readability and puts puts three levels of asserts on the function:
      first level is the number of input tensors and output tensors; second level is
      the rank of each tensor; third level is the size of each dimension of each
      tensor. The decorator inserts those asserts right before and right after
      'forward' call.
  
      First level, assert on number of inputs and outputs. In the representation
      input tensors are separated from output tensors by an arrow '->'. For layers
      taking multiple input tensors or returning multiple output tensors, those
      tensors will be separated by a comma ','.
      For example, specification 'bsd,df->bsf' asserts that there will be two
      input tensors, with shapes represented by 'bsd' and 'df' respectively; and
      a single output tensor with shape represented by 'bsf'.
  
      Second level, asserts on possible rank of each tensor. Most commonly,
      each letter represents a single dimension. For example,the tensor with shapes
      represented by 'bsd' has rank three; with 'df' it has rank two. The special
      case is an ellipsis ('...'), which expand to arbitrary number of dimensions,
      including zero. For example, the tensor with specification '...sf' has at
      least two dimensions. Each tensor may have in its representation one ellipsis.
  
      Third level, asserts the size of each dimension. If two dimensions in any
      of input or output tensors have the same letter in the representation then
      they must have the same size. For example, with a tensor A represented by 'df'
      and a tensor B represented by 'bsf', the size of the second dimension of A
      must equal the size of the third dimension of B. Another example: with a
      tensor C represented by '...dv' and a tensor D represented by 'd', the size of
      the first and only dimension of D must be equal to the size of the second to
      last dimension of tensor C.
  
      If two distinct tensors have an ellipsis in their representation then all of
      dimensions covered by those ellipses must match. For example, with a tensor E
      represented by '...d' and tensor F represented by '...x' then E and F must
      have the same rank, and the sizes of all but the last dimensions must match.
  
      Examples:
      # In Dense layer there is a single input and single output; the last dimension
      # may change in size, while the sizes of all previous dimensions, marked by
      # an ellipsis, will stay the same.
      @assert_shape('...a->...b')
      class Dense(base.Layer):
        (...)
  
      # DotProductCausalAttention takes three tensors as input: Queries, Keys, and
      # Values, and outputs a single tensor. Sizes of the first two dimensions in
      # all those tensors must match, while the last dimension must match only
      # between Queries and Keys, and separately between Values and output tensor.
      @assert_shape('blk,blk,bld->bld')
      class DotProductCausalAttention(base.Layer):
        (...)
  
      # assert_shape can also be placed before the function returning base.Layer.
      @assert_shape('...d->...')
      def ReduceSum():
        return Fn('ReduceSum', lambda x: jnp.sum(x, axis=-1, keepdims=False))
  
      Args:
        specification: A text specification for the input/output tensors.
  
      Returns:
        The decorator changing the class or function.
      """
      caller = inspect.getframeinfo(inspect.stack()[1][0])
      message = f"Defined at {caller.filename}:{caller.lineno}"
  
      def wrap_cls(cls):
          forward = getattr(cls, "forward")
          init = getattr(cls, "__init__")
  
          before_spec, after_spec = specification.split("->")
  
          @functools.wraps(init)
          def init_wrapper(self, *args, **kwargs):
              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
              out = init(self, *args, **kwargs)
              self._before_assert_fun = (
                  before_assert  # pylint: disable=protected-access
              )
              self._after_assert_fun = (
                  after_assert  # pylint: disable=protected-access
              )
              return out
  
          @functools.wraps(forward)
          def forward_wrapper(self, x, *args, **kwargs):
              x = self._before_assert_fun.forward(
                  x
              )  # pylint: disable=protected-access
              y = forward(self, x, *args, **kwargs)
              y = self._after_assert_fun.forward(
                  y
              )  # pylint: disable=protected-access
              return y
  
          setattr(cls, "forward", forward_wrapper)
          setattr(cls, "__init__", init_wrapper)
          return cls
  
      # TODO(jaszczur): replace this with forward/init override.
      def wrap_fun(fun):
          @functools.wraps(fun)
          def fun_wrapper(*args, **kwargs):
              layer = fun(*args, **kwargs)
              return AssertFunction(specification, layer, message)
  
          return fun_wrapper
  
      def wrap_fun_or_cls(fun_or_cls):
          return (
              wrap_cls(fun_or_cls)
              if inspect.isclass(fun_or_cls)
              else wrap_fun(fun_or_cls)
          )
  
      return wrap_fun_or_cls
  
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 1 for calling a global function with example usage: assert_shape(specification) and returns: wrap_fun_or_cls, cls, fun_wrapper, unknown, out, y, AssertFunction