def PrintShape(n_in=1, msg=''):
  """Prints the shapes of `n_in` inputs and returns then unchanged."""
  def Fwd(xs):
    def format_shape(x):  # pylint: disable = invalid-name
      return str(x.shape) + f'[{x.dtype}]'
    if n_in > 1:
      shapes_and_dtypes = ', '.join([format_shape(x) for x in xs])
    else:
      shapes_and_dtypes = format_shape(xs)
    info = f'PrintShape: {msg}: [{shapes_and_dtypes}]'
    print(info)
    logging.info(info)
    return xs
  return base.PureLayer(Fwd, n_in=n_in, n_out=n_in, name=f'PrintShape_{n_in}')
Prints the shapes of `n_in` inputs and returns then unchanged.
trax/layers/core.py