import tensorflow as tf

      ]
    else:
      return self.create_accumulator()

  def extract_output(self, accumulator):
    # For each output, cast that output to the specified type. Note there
    # will be one output for each input tensor to the analyzer.
    return [
        sub_accumulator.astype(output_dtype) for sub_accumulator, output_dtype
        in zip(accumulator, self._output_dtypes)
    ]

  def output_tensor_infos(self):
    return [
        analyzer_nodes.TensorInfo(tf.as_dtype(dtype), shape, None)
        for dtype, shape in zip(self._output_dtypes, self._output_shapes)
    ]


def _get_output_shape_from_input(x):
  if isinstance(x, tf.SparseTensor):
    return x.get_shape().as_list()[1:]

  # When reducing over batch dimensions, with known shape, the result will be
  # the same shape as the input, but without the batch.
  if x.shape.rank is not None:
    return x.shape.as_list()[1:]
  return (None,)
