class_docstr: The Batch-Reshaping distribution.
  
  This "meta-distribution" reshapes the batch dimensions of another
  distribution.
  
  #### Examples
  
  ```python
  tfd = tfp.distributions
  
  dtype = np.float32
  dims = 2
  new_batch_shape = [1, 2, -1]
  old_batch_shape = [6]
  
  scale = np.ones(old_batch_shape + [dims], dtype)
  mvn = tfd.MultivariateNormalDiag(scale_diag=scale)
  reshape_mvn = tfd.BatchReshape(
      distribution=mvn,
      batch_shape=new_batch_shape,
      validate_args=True)
  
  reshape_mvn.batch_shape
  # ==> [1, 2, 3]
  
  x = reshape_mvn.sample(sample_shape=[4, 5])
  x.shape
  # ==> [4, 5, 1, 2, 3, 2] == sample_shape + new_batch_shape + [dims]
  
  reshape_mvn.log_prob(x).shape
  # ==> [4, 5, 1, 2, 3] == sample_shape + new_batch_shape
  ```
class_docstr_tok: ["The","Batch-Reshaping","distribution",".","This","``","meta-distribution","''","reshapes","the","batch","dimensions","of","another","distribution",".","#","#","#","#","Examples","``","`","python","tfd","=","tfp.distributions","dtype","=","np.float32","dims","=","2","new_batch_shape","=","[","1",",","2",",","-1","]","old_batch_shape","=","[","6","]","scale","=","np.ones","(","old_batch_shape","+","[","dims","]",",","dtype",")","mvn","=","tfd.MultivariateNormalDiag","(","scale_diag=scale",")","reshape_mvn","=","tfd.BatchReshape","(","distribution=mvn",",","batch_shape=new_batch_shape",",","validate_args=True",")","reshape_mvn.batch_shape","#","==",">","[","1",",","2",",","3","]","x","=","reshape_mvn.sample","(","sample_shape=","[","4",",","5","]",")","x.shape","#","==",">","[","4",",","5",",","1",",","2",",","3",",","2","]","==","sample_shape","+","new_batch_shape","+","[","dims","]","reshape_mvn.log_prob","(","x",")",".shape","#","==",">","[","4",",","5",",","1",",","2",",","3","]","==","sample_shape","+","new_batch_shape","``","`"]
docstr: Helper which validates sample arg, e.g., input to `log_prob`.
docstr_tok: ["Helper","which","validates","sample","arg",",","e.g.",",","input","to","`","log_prob","`","."]
code: def _sample_control_dependencies(self, x):
      """Helper which validates sample arg, e.g., input to `log_prob`."""
      x_ndims = (
          tf.rank(x)
          if tensorshape_util.rank(x.shape) is None
          else tensorshape_util.rank(x.shape)
      )
      event_ndims = (
          tf.size(self.event_shape_tensor())
          if tensorshape_util.rank(self.event_shape) is None
          else tensorshape_util.rank(self.event_shape)
      )
      batch_ndims = (
          tf.size(self._batch_shape_unexpanded)
          if tensorshape_util.rank(self.batch_shape) is None
          else tensorshape_util.rank(self.batch_shape)
      )
      expected_batch_event_ndims = batch_ndims + event_ndims
  
      if isinstance(x_ndims, int) and isinstance(
          expected_batch_event_ndims, int
      ):
          if x_ndims < expected_batch_event_ndims:
              raise NotImplementedError(
                  "Broadcasting is not supported; too few batch and event dims "
                  "(expected at least {}, saw {}).".format(
                      expected_batch_event_ndims, x_ndims
                  )
              )
          ndims_assertion = []
      elif self.validate_args:
          ndims_assertion = [
              assert_util.assert_greater_equal(
                  x_ndims,
                  expected_batch_event_ndims,
                  message=(
                      "Broadcasting is not supported; too few "
                      "batch and event dims."
                  ),
                  name="assert_batch_and_event_ndims_large_enough",
              ),
          ]
  
      if tensorshape_util.is_fully_defined(
          self.batch_shape
      ) and tensorshape_util.is_fully_defined(self.event_shape):
          expected_batch_event_shape = np.int32(
              tensorshape_util.concatenate(
                  self.batch_shape, self.event_shape
              )
          )
      else:
          expected_batch_event_shape = tf.concat(
              [
                  self.batch_shape_tensor(),
                  self.event_shape_tensor(),
              ],
              axis=0,
          )
  
      sample_ndims = x_ndims - expected_batch_event_ndims
      if isinstance(sample_ndims, int):
          sample_ndims = max(sample_ndims, 0)
      if isinstance(
          sample_ndims, int
      ) and tensorshape_util.is_fully_defined(x.shape[sample_ndims:]):
          actual_batch_event_shape = np.int32(x.shape[sample_ndims:])
      else:
          sample_ndims = tf.maximum(sample_ndims, 0)
          actual_batch_event_shape = tf.shape(x)[sample_ndims:]
  
      assertions = []
      if isinstance(
          expected_batch_event_shape, np.ndarray
      ) and isinstance(actual_batch_event_shape, np.ndarray):
          if any(
              expected_batch_event_shape != actual_batch_event_shape
          ):
              raise NotImplementedError(
                  "Broadcasting is not supported; "
                  "unexpected batch and event shape "
                  "(expected {}, saw {}).".format(
                      expected_batch_event_shape,
                      actual_batch_event_shape,
                  )
              )
          # We need to set the final runtime-assertions to `ndims_assertion` since
          # its possible this assertion was created. We could add a condition to
          # only do so if `self.validate_args == True`, however this is redundant
          # as `ndims_assertion` already encodes this information.
          assertions.extend(ndims_assertion)
      elif self.validate_args:
          # We need to make the `ndims_assertion` a control dep because otherwise
          # TF itself might raise an exception owing to this assertion being
          # ill-defined, ie, one cannot even compare different rank Tensors.
          with tf.control_dependencies(ndims_assertion):
              shape_assertion = assert_util.assert_equal(
                  expected_batch_event_shape,
                  actual_batch_event_shape,
                  message=(
                      "Broadcasting is not supported; "
                      "unexpected batch and event shape."
                  ),
                  name="assert_batch_and_event_shape_same",
              )
          assertions.append(shape_assertion)
  
      return assertions
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import auto_composite_tensor
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 29 for calling the _BatchReshape obj's underlying member method code with example usage: obj._sample_control_dependencies(x) and returns: assertions