original_docstring: Use SoftmaxNormal quantiles to form quadrature on `K - 1` simplex.
  
    Args:
      loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
        Represents the `location` parameter of the SoftmaxNormal used for
        selecting one of the `K` affine transformations.
      scale: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
        Represents the `scale` parameter of the SoftmaxNormal used for
        selecting one of the `K` affine transformations.
      quadrature_size: Python scalar `int` representing the number of quadrature
        points.
      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.
      name: Python `str` name prefixed to Ops created by this class.
  
    Returns:
      grid: Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
        convex combination of affine parameters for `K` components.
        `grid[..., :, n]` is the `n`-th grid point, living in the `K - 1` simplex.
      probs:  Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
        associated with each grid point.
    
docstring: Use SoftmaxNormal quantiles to form quadrature on `K - 1` simplex.
docstring_tokens: Use SoftmaxNormal quantiles to form quadrature on ` K - 1 ` simplex .
code: def quadrature_scheme_softmaxnormal_quantiles(
      loc, scale, quadrature_size,
      validate_args=False, name=None):
    """Use SoftmaxNormal quantiles to form quadrature on `K - 1` simplex.
  
    Args:
      loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
        Represents the `location` parameter of the SoftmaxNormal used for
        selecting one of the `K` affine transformations.
      scale: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
        Represents the `scale` parameter of the SoftmaxNormal used for
        selecting one of the `K` affine transformations.
      quadrature_size: Python scalar `int` representing the number of quadrature
        points.
      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.
      name: Python `str` name prefixed to Ops created by this class.
  
    Returns:
      grid: Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
        convex combination of affine parameters for `K` components.
        `grid[..., :, n]` is the `n`-th grid point, living in the `K - 1` simplex.
      probs:  Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
        associated with each grid point.
    """
    with ops.name_scope(name, "softmax_normal_grid_and_probs", [loc, scale]):
      loc = ops.convert_to_tensor(loc, name="loc")
      dt = loc.dtype.base_dtype
      scale = ops.convert_to_tensor(scale, dtype=dt, name="scale")
  
      loc = maybe_check_quadrature_param(loc, "loc", validate_args)
      scale = maybe_check_quadrature_param(scale, "scale", validate_args)
  
      dist = normal_lib.Normal(loc=loc, scale=scale)
  
      def _get_batch_ndims():
        """Helper to get dist.batch_shape.ndims, statically if possible."""
        ndims = dist.batch_shape.ndims
        if ndims is None:
          ndims = array_ops.shape(dist.batch_shape_tensor())[0]
        return ndims
      batch_ndims = _get_batch_ndims()
  
      def _get_final_shape(qs):
        """Helper to build `TensorShape`."""
        bs = dist.batch_shape.with_rank_at_least(1)
        num_components = bs[-1].value
        if num_components is not None:
          num_components += 1
        tail = tensor_shape.TensorShape([num_components, qs])
        return bs[:-1].concatenate(tail)
  
      def _compute_quantiles():
        """Helper to build quantiles."""
        # Omit {0, 1} since they might lead to Inf/NaN.
        zero = array_ops.zeros([], dtype=dist.dtype)
        edges = math_ops.linspace(zero, 1., quadrature_size + 3)[1:-1]
        # Expand edges so its broadcast across batch dims.
        edges = array_ops.reshape(edges, shape=array_ops.concat([
            [-1], array_ops.ones([batch_ndims], dtype=dtypes.int32)], axis=0))
        quantiles = dist.quantile(edges)
        quantiles = SoftmaxCentered(event_ndims=1).forward(quantiles)
        # Cyclically permute left by one.
        perm = array_ops.concat([
            math_ops.range(1, 1 + batch_ndims), [0]], axis=0)
        quantiles = array_ops.transpose(quantiles, perm)
        quantiles.set_shape(_get_final_shape(quadrature_size + 1))
        return quantiles
      quantiles = _compute_quantiles()
  
      # Compute grid as quantile midpoints.
      grid = (quantiles[..., :-1] + quantiles[..., 1:]) / 2.
      # Set shape hints.
      grid.set_shape(_get_final_shape(quadrature_size))
  
      # By construction probs is constant, i.e., `1 / quadrature_size`. This is
      # important, because non-constant probs leads to non-reparameterizable
      # samples.
      probs = array_ops.fill(
          dims=[quadrature_size],
          value=1. / math_ops.cast(quadrature_size, dist.dtype))
  
      return grid, probs
code_tokens: def quadrature_scheme_softmaxnormal_quantiles ( loc , scale , quadrature_size , validate_args = False , name = None ) : with ops . name_scope ( name , "softmax_normal_grid_and_probs" , [ loc , scale ] ) : loc = ops . convert_to_tensor ( loc , name = "loc" ) dt = loc . dtype . base_dtype scale = ops . convert_to_tensor ( scale , dtype = dt , name = "scale" ) loc = maybe_check_quadrature_param ( loc , "loc" , validate_args ) scale = maybe_check_quadrature_param ( scale , "scale" , validate_args ) dist = normal_lib . Normal ( loc = loc , scale = scale ) def _get_batch_ndims ( ) : """Helper to get dist.batch_shape.ndims, statically if possible.""" ndims = dist . batch_shape . ndims if ndims is None : ndims = array_ops . shape ( dist . batch_shape_tensor ( ) ) [ 0 ] return ndims batch_ndims = _get_batch_ndims ( ) def _get_final_shape ( qs ) : """Helper to build `TensorShape`.""" bs = dist . batch_shape . with_rank_at_least ( 1 ) num_components = bs [ - 1 ] . value if num_components is not None : num_components += 1 tail = tensor_shape . TensorShape ( [ num_components , qs ] ) return bs [ : - 1 ] . concatenate ( tail ) def _compute_quantiles ( ) : """Helper to build quantiles.""" zero = array_ops . zeros ( [ ] , dtype = dist . dtype ) edges = math_ops . linspace ( zero , 1. , quadrature_size + 3 ) [ 1 : - 1 ] edges = array_ops . reshape ( edges , shape = array_ops . concat ( [ [ - 1 ] , array_ops . ones ( [ batch_ndims ] , dtype = dtypes . int32 ) ] , axis = 0 ) ) quantiles = dist . quantile ( edges ) quantiles = SoftmaxCentered ( event_ndims = 1 ) . forward ( quantiles ) perm = array_ops . concat ( [ math_ops . range ( 1 , 1 + batch_ndims ) , [ 0 ] ] , axis = 0 ) quantiles = array_ops . transpose ( quantiles , perm ) quantiles . set_shape ( _get_final_shape ( quadrature_size + 1 ) ) return quantiles quantiles = _compute_quantiles ( ) grid = ( quantiles [ ... , : - 1 ] + quantiles [ ... , 1 : ] ) / 2. grid . set_shape ( _get_final_shape ( quadrature_size ) ) probs = array_ops . fill ( dims = [ quadrature_size ] , value = 1. / math_ops . cast ( quadrature_size , dist . dtype ) ) return grid , probs
short_docstring: Use SoftmaxNormal quantiles to form quadrature on `K - 1` simplex.
short_docstring_tokens: Use SoftmaxNormal quantiles to form quadrature on ` K - 1 ` simplex .
comment: """Use SoftmaxNormal quantiles to form quadrature on `K - 1` simplex.
  
    Args:
      loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
        Represents the `location` parameter of the SoftmaxNormal used for
        selecting one of the `K` affine transformations.
      scale: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
        Represents the `scale` parameter of the SoftmaxNormal used for
        selecting one of the `K` affine transformations.
      quadrature_size: Python scalar `int` representing the number of quadrature
        points.
      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.
      name: Python `str` name prefixed to Ops created by this class.
  
    Returns:
      grid: Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
        convex combination of affine parameters for `K` components.
        `grid[..., :, n]` is the `n`-th grid point, living in the `K - 1` simplex.
      probs:  Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
        associated with each grid point.
    """$SEP$"""Helper to get dist.batch_shape.ndims, statically if possible."""$SEP$"""Helper to build `TensorShape`."""$SEP$"""Helper to build quantiles."""$SEP$# Omit {0, 1} since they might lead to Inf/NaN.$SEP$# Expand edges so its broadcast across batch dims.$SEP$# Cyclically permute left by one.$SEP$# Compute grid as quantile midpoints.$SEP$# Set shape hints.$SEP$# By construction probs is constant, i.e., `1 / quadrature_size`. This is$SEP$# important, because non-constant probs leads to non-reparameterizable$SEP$# samples.