class_docstr: Layer which put asserts on shapes of tensors, and returns them unchanged.
  
  It borrows the notation from assert_shape decorator, except it doesn't have
  the arrow '->' special character, as the input tensors are the same as output.
class_docstr_tok: ["Layer","which","put","asserts","on","shapes","of","tensors",",","and","returns","them","unchanged",".","It","borrows","the","notation","from","assert_shape","decorator",",","except","it","does","n't","have","the","arrow","'-",">","'","special","character",",","as","the","input","tensors","are","the","same","as","output","."]
code: def forward(self, xs):
      if not self._linked:
          for k in list(self._defined_shapes.keys()):
              if not k.isdigit():
                  del self._defined_shapes[k]
  
      if not isinstance(xs, (list, tuple)):
          xs = [xs]
  
      # Try-except below checks if something is wrong with shapes. It can happen
      # e.g. when using trax2keras. If this is the case we cannot check if shapes
      # are correct or not
      try:
          for x in xs:
              for i in range(len(x.shape)):
                  if x.shape[i] != x.shape[i]:
                      raise TypeError()
      except TypeError:
          message = (
              "AssertShape cannot check shapes. This often happens when"
              " using trax2keras. Shape asserts are skipped."
          )
          print(message)
          logging.warning(message)
          if len(xs) == 1:
              return xs[0]
          else:
              return xs
  
      # helper functions
      def assert_true(cond):
          if not cond:
              shapes = [x.shape for x in xs]
              defined_shapes_dict_without_digits = {
                  k: v
                  for k, v in self._defined_shapes.items()
                  if not k.isdigit()
              }
              raise ValueError(
                  f"AssertShape Error. Expected {self._specs}, got {shapes} with dict"
                  f" {defined_shapes_dict_without_digits}. {self._message}"
              )
  
      def assert_equal(a, b):
          assert_true(a == b)
          return a
  
      def check_shape(shape, spec):
          assert_equal(len(shape), len(spec))
          for shape_dim, letter in zip(shape, spec):
              if letter in self._defined_shapes:
                  self._defined_shapes[letter] = assert_equal(
                      self._defined_shapes[letter], shape_dim
                  )
              else:
                  self._defined_shapes[letter] = shape_dim
  
      def check_ellipsys(shape):
          if "." not in self._defined_shapes:
              self._defined_shapes["."] = shape
          else:
              assert_equal(len(shape), len(self._defined_shapes["."]))
              for s1, s2 in zip(shape, self._defined_shapes["."]):
                  assert_equal(s1, s2)
  
      # actual asserts
      assert_equal(len(xs), len(self._specs))
  
      for x, spec in zip(xs, self._specs):
          if "." in spec:
              assert_true(len(x.shape) >= (len(spec) - 1))
  
              before, after = spec.split(".")
              check_shape(x.shape[: len(before)], before)
              if after:
                  check_shape(x.shape[-len(after) :], after)
                  check_ellipsys(x.shape[len(before) : -len(after)])
              else:
                  # if len(after) == 0 then -len(after) in indices evaluates badly.
                  check_ellipsys(x.shape[len(before) :])
          else:
              check_shape(x.shape, spec)
  
      if len(xs) == 1:
          return xs[0]
      else:
          return xs
  
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 the AssertShape obj's underlying member method code with example usage: obj.forward(xs) and returns: a, xs, xs, xs, xs