class_docstr: The Batch-Concatenating distribution.
  
  This distribution concatenates a list of distributions
  along a given axis of their batch shapes.
  
  In order to be concatenable along `axis`,
  `dist1`  and `dist2` should meet the following requirements:
  1. `dist1.event_shape == dist2.event_shape`
  2. `len(dist1.batch_shape) == len(dist2.batch_shape)`
  3. `\forall i \neq axis, dist1.batch_shape[i] == 1 or
                           dist1.batch_shape[i] == dist2.batch_shape[i]`
  If `dist1.batch_shape[i] == 1`, the distribution will be
  broadcasted along this axis.
  
  TODO(b/179916710): The distribution does NOT support sample broadcasting.
  Given a sample `x` and a BatchConcat distribution `dist`, `dist.log_prob(x)`
  requires that `x.shape == sample_shape + dist.batch_shape + dist.event_shape`.
  
  #### Example
  
  ```python
  tfd = tfp.distributions
  
  dtype = np.float32
  dims = 2
  batch_shape_1 = [32, 2]
  batch_shape_2 = [32, 6]
  
  scale_1 = np.ones(batch_shape_1 + [dims], dtype)
  mvn_1 = tfd.MultivariateNormalDiag(scale_diag=scale_1)
  scale_2 = np.ones(batch_shape_2 + [dims], dtype)
  mvn_2 = tfd.MultivariateNormalDiag(scale_diag=scale_2)
  batched_mvn = tfd.BatchConcat(
      distributions=[mvn_1, mvn_2],
      axis=1,
      validate_args=True)
  
  batched_mvn.batch_shape
  # ==> [32, 8]
  
  x = batched_mvn.sample(sample_shape=[4, 5])
  x.shape
  # ==> [4, 5, 32, 8, 2] == sample_shape + batched_mvn.batch_shape + [dims]
  
  reshape_mvn.log_prob(x).shape
  # ==> [4, 5, 32, 8] == sample_shape + batched_mvn.batch_shape
  ```
  
  #### Example for broadcasting.
  The distributions need to have the same batch shape size
  to enable broadcasting.
  
  ```python
  tfd = tfp.distributions
  
  dtype = np.float32
  dims = 2
  batch_shape_1 = [1, 2, 16]
  batch_shape_2 = [32, 6, 1]
  
  scale_1 = np.ones(batch_shape_1 + [dims], dtype)
  mvn_1 = tfd.MultivariateNormalDiag(scale_diag=scale_1)
  scale_2 = np.ones(batch_shape_2 + [dims], dtype)
  mvn_2 = tfd.MultivariateNormalDiag(scale_diag=scale_2)
  batched_mvn = tfd.BatchConcat(
      distributions=[mvn_1, mvn_2],
      axis=1,
      validate_args=True)
  
  batched_mvn.batch_shape
  # ==> [32, 8, 16]
  
  x = batched_mvn.sample(sample_shape=[4, 5])
  x.shape
  # ==> [4, 5, 32, 8, 16, 2] == sample_shape + batched_mvn.batch_shape + [dims]
  
  reshape_mvn.log_prob(x).shape
  # ==> [4, 5, 32, 8, 16] == sample_shape + batched_mvn.batch_shape
  ```
class_docstr_tok: ["The","Batch-Concatenating","distribution",".","This","distribution","concatenates","a","list","of","distributions","along","a","given","axis","of","their","batch","shapes",".","In","order","to","be","concatenable","along","`","axis","`",",","`","dist1","`","and","`","dist2","`","should","meet","the","following","requirements",":","1",".","`","dist1.event_shape","==","dist2.event_shape","`","2",".","`","len","(","dist1.batch_shape",")","==","len","(","dist2.batch_shape",")","`","3",".","`","\\forall","i","\\neq","axis",",","dist1.batch_shape","[","i","]","==","1","or","dist1.batch_shape","[","i","]","==","dist2.batch_shape","[","i","]","`","If","`","dist1.batch_shape","[","i","]","==","1","`",",","the","distribution","will","be","broadcasted","along","this","axis",".","TODO","(","b\/179916710",")",":","The","distribution","does","NOT","support","sample","broadcasting",".","Given","a","sample","`","x","`","and","a","BatchConcat","distribution","`","dist","`",",","`","dist.log_prob","(","x",")","`","requires","that","`","x.shape","==","sample_shape","+","dist.batch_shape","+","dist.event_shape","`",".","#","#","#","#","Example","``","`","python","tfd","=","tfp.distributions","dtype","=","np.float32","dims","=","2","batch_shape_1","=","[","32",",","2","]","batch_shape_2","=","[","32",",","6","]","scale_1","=","np.ones","(","batch_shape_1","+","[","dims","]",",","dtype",")","mvn_1","=","tfd.MultivariateNormalDiag","(","scale_diag=scale_1",")","scale_2","=","np.ones","(","batch_shape_2","+","[","dims","]",",","dtype",")","mvn_2","=","tfd.MultivariateNormalDiag","(","scale_diag=scale_2",")","batched_mvn","=","tfd.BatchConcat","(","distributions=","[","mvn_1",",","mvn_2","]",",","axis=1",",","validate_args=True",")","batched_mvn.batch_shape","#","==",">","[","32",",","8","]","x","=","batched_mvn.sample","(","sample_shape=","[","4",",","5","]",")","x.shape","#","==",">","[","4",",","5",",","32",",","8",",","2","]","==","sample_shape","+","batched_mvn.batch_shape","+","[","dims","]","reshape_mvn.log_prob","(","x",")",".shape","#","==",">","[","4",",","5",",","32",",","8","]","==","sample_shape","+","batched_mvn.batch_shape","``","`","#","#","#","#","Example","for","broadcasting",".","The","distributions","need","to","have","the","same","batch","shape","size","to","enable","broadcasting",".","``","`","python","tfd","=","tfp.distributions","dtype","=","np.float32","dims","=","2","batch_shape_1","=","[","1",",","2",",","16","]","batch_shape_2","=","[","32",",","6",",","1","]","scale_1","=","np.ones","(","batch_shape_1","+","[","dims","]",",","dtype",")","mvn_1","=","tfd.MultivariateNormalDiag","(","scale_diag=scale_1",")","scale_2","=","np.ones","(","batch_shape_2","+","[","dims","]",",","dtype",")","mvn_2","=","tfd.MultivariateNormalDiag","(","scale_diag=scale_2",")","batched_mvn","=","tfd.BatchConcat","(","distributions=","[","mvn_1",",","mvn_2","]",",","axis=1",",","validate_args=True",")","batched_mvn.batch_shape","#","==",">","[","32",",","8",",","16","]","x","=","batched_mvn.sample","(","sample_shape=","[","4",",","5","]",")","x.shape","#","==",">","[","4",",","5",",","32",",","8",",","16",",","2","]","==","sample_shape","+","batched_mvn.batch_shape","+","[","dims","]","reshape_mvn.log_prob","(","x",")",".shape","#","==",">","[","4",",","5",",","32",",","8",",","16","]","==","sample_shape","+","batched_mvn.batch_shape","``","`"]
docstr: Construct BatchConcat distribution.
  
  Args:
    distributions: A sequence of `Distribution` instances to concatenate.
      Instances of the distributions should have the same event_shape and
      compatible batch_shapes that can be concatenated along the specified
      axis.
    axis: Positive `int` specifying the axis along which to concatenate
      the given distributions.
    validate_args: Python `bool`, default `False`. When `True` distribution
      parameters are checked for validity despite possibly degrading runtime
      performance. When `False` invalid inputs may silently render incorrect
      outputs.
    allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
      (e.g., mean, mode, variance) use the value `NaN` to indicate the
      result is undefined. When `False`, an exception is raised if one or
      more of the statistic's batch members are undefined.
    name: The name to give Ops created by this `Distribution`.
      Default value: "BatchConcat"
  
  Raises:
    ValueError: if `event_shape` of given distributions are not the same.
    ValueError: if `batch_shape` of given distributions are not concatenable.
    ValueError: if distributions don't have the same type.
    ValueError: if axis is out of range or negative.
docstr_tok: ["Construct","BatchConcat","distribution",".","Args",":","distributions",":","A","sequence","of","`","Distribution","`","instances","to","concatenate",".","Instances","of","the","distributions","should","have","the","same","event_shape","and","compatible","batch_shapes","that","can","be","concatenated","along","the","specified","axis",".","axis",":","Positive","`","int","`","specifying","the","axis","along","which","to","concatenate","the","given","distributions",".","validate_args",":","Python","`","bool","`",",","default","`","False","`",".","When","`","True","`","distribution","parameters","are","checked","for","validity","despite","possibly","degrading","runtime","performance",".","When","`","False","`","invalid","inputs","may","silently","render","incorrect","outputs",".","allow_nan_stats",":","Python","`","bool","`",",","default","`","True","`",".","When","`","True","`",",","statistics","(","e.g.",",","mean",",","mode",",","variance",")","use","the","value","`","NaN","`","to","indicate","the","result","is","undefined",".","When","`","False","`",",","an","exception","is","raised","if","one","or","more","of","the","statistic","'s","batch","members","are","undefined",".","name",":","The","name","to","give","Ops","created","by","this","`","Distribution","`",".","Default","value",":","``","BatchConcat","''","Raises",":","ValueError",":","if","`","event_shape","`","of","given","distributions","are","not","the","same",".","ValueError",":","if","`","batch_shape","`","of","given","distributions","are","not","concatenable",".","ValueError",":","if","distributions","do","n't","have","the","same","type",".","ValueError",":","if","axis","is","out","of","range","or","negative","."]
code: def __init__(
      self,
      distributions,
      axis,
      validate_args=False,
      allow_nan_stats=True,
      name="BatchConcat",
  ):
      """Construct BatchConcat distribution.
  
      Args:
        distributions: A sequence of `Distribution` instances to concatenate.
          Instances of the distributions should have the same event_shape and
          compatible batch_shapes that can be concatenated along the specified
          axis.
        axis: Positive `int` specifying the axis along which to concatenate
          the given distributions.
        validate_args: Python `bool`, default `False`. When `True` distribution
          parameters are checked for validity despite possibly degrading runtime
          performance. When `False` invalid inputs may silently render incorrect
          outputs.
        allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
          (e.g., mean, mode, variance) use the value `NaN` to indicate the
          result is undefined. When `False`, an exception is raised if one or
          more of the statistic's batch members are undefined.
        name: The name to give Ops created by this `Distribution`.
          Default value: "BatchConcat"
  
      Raises:
        ValueError: if `event_shape` of given distributions are not the same.
        ValueError: if `batch_shape` of given distributions are not concatenable.
        ValueError: if distributions don't have the same type.
        ValueError: if axis is out of range or negative.
      """
      parameters = dict(locals())
      with tf.name_scope(name) as name:
          self._distributions = distributions
          self._axis = axis
  
          dtype = dtype_util.common_dtype(distributions)
  
          reparameterizable = all(
              d._reparameterization_type
              == reparameterization.FULLY_REPARAMETERIZED
              for d in self._distributions
          )
  
          reparameterization_type = (
              reparameterization.FULLY_REPARAMETERIZED
              if reparameterizable
              else reparameterization.NOT_REPARAMETERIZED
          )
  
          super(_BatchConcat, self).__init__(
              dtype=dtype,
              reparameterization_type=reparameterization_type,
              validate_args=validate_args,
              allow_nan_stats=allow_nan_stats,
              parameters=parameters,
              name=name,
          )
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  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 reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 to create a new _BatchConcat object from inherited base classes: distribution_lib with example: obj = _BatchConcat(distributions, axis, validate_args, allow_nan_stats, name)