code: 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}"
      )
  
imports: from absl import logging
  import numpy
  import tensorflow
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
label_desc: Use image node_id 1 for calling a global function with example usage: PrintShape(n_in, msg) and returns: base, xs, unknown