code: def test_docstring_categorical(self):
      y = tf.constant([[0, 1], [2, 1], [1, 1], [2, 2]], dtype=tf.int64)
      logits = [
          [[0.0, -1, -2], [-1.0, 0, -1]],
          [[2.0, -1, 0], [-1.0, 0, -1]],
          [[-1.0, 0, -1], [0.0, -2, -1]],
          [[-1.0, 0, -2], [0, -1, -2]],
      ]
      d = categorical.Categorical(logits=logits)
  
      def all_categories(d):
          num_classes = tf.shape(d.logits_parameter())[-1]
          batch_ndims = tf.size(d.batch_shape_tensor())
          expand_shape = tf.pad(
              [num_classes],
              paddings=[[0, batch_ndims]],
              constant_values=1,
          )
          return tf.reshape(
              tf.range(num_classes, dtype=d.dtype), expand_shape
          )
  
      all_pred_log_prob = d.log_prob(all_categories(d))
      yhat = tf.argmax(all_pred_log_prob, axis=0)
  
      def rollaxis(x, shift):
          return tf.transpose(
              x, tf.roll(tf.range(tf.rank(x)), shift=shift, axis=0)
          )
  
      pred_log_prob = tf.gather(
          rollaxis(all_pred_log_prob, shift=-1),
          yhat,
          batch_dims=len(d.batch_shape),
      )
      hit = tf.equal(y, yhat)
      ece, acc, conf, cnt, edges, buckets = self.evaluate(
          calibration.expected_calibration_error_quantiles(
              hit, pred_log_prob, num_buckets=10
          )
      )
      self.assertEqual((2,), ece.shape)
      self.assertEqual((10, 2), acc.shape)
      self.assertEqual((10, 2), conf.shape)
      self.assertEqual((10, 2), cnt.shape)
      self.assertEqual((11, 2), edges.shape)
      self.assertEqual((4, 2), buckets.shape)
  
imports: import itertools
  from absl import logging
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.stats import calibration
label_desc: Use image node_id 2 for calling the ExpectedCalibrationErrorQuantiles obj's underlying member method code with example usage: obj.test_docstring_categorical() and returns: tf, tf