tensorflow_probability/python/internal/distribution_util.py
def pad_mixture_dimensions(x, mixture_distribution, categorical_distribution,
                           event_ndims):
  """Pad dimensions of event tensors for mixture distributions.

  See `Mixture._sample_n` and `MixtureSameFamily._sample_n` for usage examples.

  Args:
    x: event tensor to pad.
    mixture_distribution: Base distribution of the mixture.
    categorical_distribution: `Categorical` distribution that mixes the base
      distribution.
    event_ndims: Integer specifying the number of event dimensions in the event
      tensor.

  Returns:
    A padded version of `x` that can broadcast with `categorical_distribution`.
  """
  with tf.name_scope("pad_mix_dims"):

    def _get_ndims(d):
      if tensorshape_util.rank(d.batch_shape) is not None:
        return tensorshape_util.rank(d.batch_shape)
      return tf.shape(input=d.batch_shape_tensor())[0]

    dist_batch_ndims = _get_ndims(mixture_distribution)
    cat_batch_ndims = _get_ndims(categorical_distribution)
    pad_ndims = tf.where(categorical_distribution.is_scalar_batch(),
                         dist_batch_ndims, dist_batch_ndims - cat_batch_ndims)
    s = tf.shape(input=x)
    x = tf.reshape(
        x,
        shape=tf.concat([
            s[:-1],
            tf.ones([pad_ndims], dtype=tf.int32),
            s[-1:],
            tf.ones([event_ndims], dtype=tf.int32),
        ],
                        axis=0))
    return x
def pad_mixture_dimensions(x, mixture_distribution, categorical_distribution,
                           event_ndims):
  with tf.name_scope("pad_mix_dims"):

    def _get_ndims(d):
      if tensorshape_util.rank(d.batch_shape) is not None:
        return tensorshape_util.rank(d.batch_shape)
      return tf.shape(input=d.batch_shape_tensor())[0]

    dist_batch_ndims = _get_ndims(mixture_distribution)
    cat_batch_ndims = _get_ndims(categorical_distribution)
    pad_ndims = tf.where(categorical_distribution.is_scalar_batch(),
                         dist_batch_ndims, dist_batch_ndims - cat_batch_ndims)
    s = tf.shape(input=x)
    x = tf.reshape(
        x,
        shape=tf.concat([
            s[:-1],
            tf.ones([pad_ndims], dtype=tf.int32),
            s[-1:],
            tf.ones([event_ndims], dtype=tf.int32),
        ],
                        axis=0))
    return x
def pad_mixture_dimensions ( x , mixture_distribution , categorical_distribution , event_ndims ) : with tf . name_scope ( "pad_mix_dims" ) : def _get_ndims ( d ) : if tensorshape_util . rank ( d . batch_shape ) is not None : return tensorshape_util . rank ( d . batch_shape ) return tf . shape ( input = d . batch_shape_tensor ( ) ) [ 0 ] dist_batch_ndims = _get_ndims ( mixture_distribution ) cat_batch_ndims = _get_ndims ( categorical_distribution ) pad_ndims = tf . where ( categorical_distribution . is_scalar_batch ( ) , dist_batch_ndims , dist_batch_ndims - cat_batch_ndims ) s = tf . shape ( input = x ) x = tf . reshape ( x , shape = tf . concat ( [ s [ : - 1 ] , tf . ones ( [ pad_ndims ] , dtype = tf . int32 ) , s [ - 1 : ] , tf . ones ( [ event_ndims ] , dtype = tf . int32 ) , ] , axis = 0 ) ) return x
Pad dimensions of event tensors for mixture distributions.

  See `Mixture._sample_n` and `MixtureSameFamily._sample_n` for usage examples.

  Args:
    x: event tensor to pad.
    mixture_distribution: Base distribution of the mixture.
    categorical_distribution: `Categorical` distribution that mixes the base
      distribution.
    event_ndims: Integer specifying the number of event dimensions in the event
      tensor.

  Returns:
    A padded version of `x` that can broadcast with `categorical_distribution`.
Pad dimensions of event tensors for mixture distributions .