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: Calls `fn`, appropriately reshaping its input `x` and output.
docstr_tok: ["Calls","`","fn","`",",","appropriately","reshaping","its","input","`","x","`","and","output","."]
code: def _call_reshape_input_output(
      self,
      fn,
      x,
      input_event_shape=None,
      output_event_shape=None,
      keep_event_dims=False,
      extra_kwargs=None,
  ):
      """Calls `fn`, appropriately reshaping its input `x` and output."""
      # Note: we take `extra_kwargs` as a dict rather than `**extra_kwargs`
      # because it is possible the user provided extra kwargs would itself
      # have `fn` and/or `x` as a key.
      if input_event_shape is None:
          static_input_event_shape, input_event_shape_tensor = (
              self.event_shape,
              self.event_shape_tensor(),
          )
      else:
          (
              static_input_event_shape,
              input_event_shape_tensor,
          ) = input_event_shape
  
      if output_event_shape is None:
          if input_event_shape is None:
              static_output_event_shape, output_event_shape_tensor = (
                  static_input_event_shape,
                  input_event_shape_tensor,
              )
          else:
              static_output_event_shape, output_event_shape_tensor = (
                  self.event_shape,
                  self.event_shape_tensor(),
              )
      else:
          (
              static_output_event_shape,
              output_event_shape_tensor,
          ) = output_event_shape
  
      sample_shape, static_sample_shape = self._sample_shape(
          x, static_input_event_shape, input_event_shape_tensor
      )
      old_shape = ps.concat(
          [
              sample_shape,
              self.distribution.batch_shape_tensor(),
              input_event_shape_tensor,
          ],
          axis=0,
      )
      x_reshape = tf.reshape(x, old_shape)
      result = (
          fn(x_reshape, **extra_kwargs)
          if extra_kwargs
          else fn(x_reshape)
      )
      new_shape = ps.concat(
          [
              sample_shape,
              self._batch_shape_unexpanded,
          ],
          axis=0,
      )
      if keep_event_dims:
          new_shape = ps.concat(
              [new_shape, output_event_shape_tensor], axis=0
          )
      result = tf.reshape(result, new_shape)
      if (
          tensorshape_util.rank(static_sample_shape) is not None
          and tensorshape_util.rank(self.batch_shape) is not None
      ):
          new_shape = tensorshape_util.concatenate(
              static_sample_shape, self.batch_shape
          )
          if keep_event_dims:
              new_shape = tensorshape_util.concatenate(
                  new_shape, static_output_event_shape
              )
          tensorshape_util.set_shape(result, new_shape)
      return result
  
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 27 for calling the _BatchReshape obj's underlying member method code with example usage: obj._call_reshape_input_output(fn, x, input_event_shape, output_event_shape, keep_event_dims, extra_kwargs) and returns: result