import tensorflow as tf


  def extract_output(self, accumulator):
    result = accumulator.GetQuantiles(self._num_quantiles).to_pylist()
    if not result:
      return [np.zeros(self._output_shape, self._bucket_numpy_dtype)]
    result = np.array(result, self._bucket_numpy_dtype)
    # Trim elementwise results if max and min should be excluded.
    if not self._include_max_and_min:
      result = result[:, 1:-1]
    return [np.reshape(result, self._output_shape)]

  def output_tensor_infos(self):
    return [
        analyzer_nodes.TensorInfo(
            tf.as_dtype(self._bucket_numpy_dtype), self._output_shape, None)
    ]

  @property
  def accumulator_coder(self):
    return _QuantilesSketchCacheCoder()


class _QuantilesSketchCacheCoder(analyzer_nodes.CacheCoder):
  """Cache coder for the quantiles accumulator."""

  def encode_cache(self, accumulator):
    # TODO(b/174549940): Consider exposing and calling
    # `QuantilesSketch::Serialize` directly.
