class_docstr: A layer receiving a tensor, and adding it to TensorBoard as an image.
  
  It takes an input and returns it unchanged. It stores this input as a state to
  be used as a metric in TensorBoard.
  It converts a tensor to a scalar by running a given aggregation function (mean
  by default). On TensorBoard, results for each device will be reported
  separately.
class_docstr_tok: ["A","layer","receiving","a","tensor",",","and","adding","it","to","TensorBoard","as","an","image",".","It","takes","an","input","and","returns","it","unchanged",".","It","stores","this","input","as","a","state","to","be","used","as","a","metric","in","TensorBoard",".","It","converts","a","tensor","to","a","scalar","by","running","a","given","aggregation","function","(","mean","by","default",")",".","On","TensorBoard",",","results","for","each","device","will","be","reported","separately","."]
docstr: Executes this layer as part of a forward pass through the model.
  
  Args:
    x: Tensor of same shape and dtype as the input signature used to
        initialize this layer.
  
  Returns:
    Tensor with previously specified shape and dtype.
docstr_tok: ["Executes","this","layer","as","part","of","a","forward","pass","through","the","model",".","Args",":","x",":","Tensor","of","same","shape","and","dtype","as","the","input","signature","used","to","initialize","this","layer",".","Returns",":","Tensor","with","previously","specified","shape","and","dtype","."]
code: def forward(self, x):
      """Executes this layer as part of a forward pass through the model.
  
      Args:
        x: Tensor of same shape and dtype as the input signature used to
            initialize this layer.
  
      Returns:
        Tensor with previously specified shape and dtype.
      """
      self.state = {}
      batch_size = x[0].shape[0]
      num_images = min(self._num_summaries, batch_size)
      for s in range(num_images):
          images = []
          for i in range(self._n_in):
              images.append(
                  self._recover_fn(x[i][s])
                  if self._recover_fn
                  else x[i][s]
              )
          self.state[self._name + str(s)] = jnp.concatenate(
              images, axis=0
          )
      return x[: self._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 2 for calling the SummaryImage obj's underlying member method code with example usage: obj.forward(x) and returns: x