docstr: Test batching support in Brier decomposition.
docstr_tok: ["Test","batching","support","in","Brier","decomposition","."]
code: def test_brier_decomposition_batching(self):
      """Test batching support in Brier decomposition."""
      tf.random.set_seed(1)
      nlabels = 4
      nsamples = 1024
      batch_shape = [3, 5]
      batch_size = np.prod(batch_shape, dtype=np.int32)
      seed_stream = test_util.test_seed_stream()
  
      logits = tf.random.normal(
          batch_shape + [nsamples, nlabels], seed=seed_stream()
      )
      labels = tf.random.uniform(
          batch_shape + [nsamples],
          maxval=nlabels,
          seed=seed_stream(),
          dtype=tf.int32,
      )
      flat_logits = tf.reshape(logits, [batch_size, nsamples, nlabels])
      flat_labels = tf.reshape(labels, [batch_size, nsamples])
  
      decomps = tf.stack(
          calibration.brier_decomposition(labels, logits)
      )
      decomps = tf.reshape(decomps, [3, batch_size])
  
      decomps_i = []
      for i in range(batch_size):
          d = calibration.brier_decomposition(
              flat_labels[i], flat_logits[i]
          )
          decomps_i.append(tf.stack(d))
      decomps, decomps_i = self.evaluate(
          [decomps, tf.stack(decomps_i, axis=-1)]
      )
      for i in range(batch_size):
          self.assertAllClose(
              decomps[:, i], decomps_i[:, i], rtol=0.20, atol=0.02
          )
  
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 CalibrationTest obj's underlying member method code with example usage: obj.test_brier_decomposition_batching() without return types