(specification)
return wrap_fun_or_cls
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.
Decorator for checking the input and output shapes of Layer.
Decorator for checking the input and output shapes of Layer .
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
def assert_shape ( specification ) : 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