def __init__(self, spec, message=None, visible_layer=False):
    """Creates AssertShape layer.

    Args:
      spec: Specification for input tensors. See assert_shape decorator for the
          full documentation.
      message: An optional message to include when assert fails. By default it
          includes the filename and line number where this function was called.
      visible_layer: If true, print this layer inside the model (default: False)
    """
    name = 'AssertShape' if visible_layer else ''
    super().__init__(name=name)
    spec = spec.replace('...', '.')
    for letter in spec:
      assert letter in string.ascii_letters + string.digits + '.' + ','
    self._specs = spec.split(',')
    self._n_in = self._n_out = len(self._specs)

    self._defined_shapes = {str(i): i for i in range(10)}
    self._linked = False

    if message is None:
      caller = inspect.getframeinfo(inspect.stack()[1][0])
      self._message = f'Defined at {caller.filename}:{caller.lineno}'
    else:
      self._message = message
Creates AssertShape layer.

    Args:
      spec: Specification for input tensors. See assert_shape decorator for the
          full documentation.
      message: An optional message to include when assert fails. By default it
          includes the filename and line number where this function was called.
      visible_layer: If true, print this layer inside the model (default: False)
    
trax/layers/assert_shape.py