@kullback_leibler.RegisterKL(Bernoulli, Bernoulli)<EOL>def _kl_bernoulli_bernoulli(a, b, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>delta_probs0 = tf.nn.softplus(-b.logits) - tf.nn.softplus(-a.logits)<EOL>delta_probs1 = tf.nn.softplus(b.logits) - tf.nn.softplus(a.logits)<EOL>return (tf.sigmoid(a.logits) * delta_probs0<EOL>+ tf.sigmoid(-a.logits) * delta_probs1)<EOL><DEDENT>

Calculate the batched KL divergence KL(a || b) with a and b Bernoulli.

    Args:
      a: instance of a Bernoulli distribution object.
      b: instance of a Bernoulli distribution object.
      name: (optional) Name to use for created operations.
        default is "kl_bernoulli_bernoulli".

    Returns:
      Batchwise KL(a || b)


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._logits, self._probs = distribution_util.get_logits_and_probs(<EOL>logits=logits,<EOL>probs=probs,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>super(Bernoulli, self).__init__(<EOL>dtype=dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._logits, self._probs],<EOL>name=name)<EOL>

Construct Bernoulli distributions.

        Args:
          logits: An N-D `Tensor` representing the log-odds of a `1` event. Each
            entry in the `Tensor` parametrizes an independent Bernoulli distribution
            where the probability of an event is sigmoid(logits). Only one of
            `logits` or `probs` should be passed in.
          probs: An N-D `Tensor` representing the probability of a `1`
            event. Each entry in the `Tensor` parameterizes an independent
            Bernoulli distribution. Only one of `logits` or `probs` should be passed
            in.
          dtype: The type of the event samples. Default: `int32`.
          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: Python `str` name prefixed to Ops created by this class.

        Raises:
          ValueError: If p and logits are passed, or if neither are passed.


return self._logits<EOL>


return self._probs<EOL>


return tf.cast(self.probs > <NUM_LIT:0.5>, self.dtype)<EOL>

Returns `1` if `prob > 0.5` and `0` otherwise.


if rate is not None:<EOL><INDENT>scale = rate<EOL><DEDENT>parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([concentration, scale],<EOL>preferred_dtype=tf.float32)<EOL>concentration = tf.convert_to_tensor(<EOL>value=concentration, name="<STR_LIT>", dtype=dtype)<EOL>scale = tf.convert_to_tensor(value=scale, name="<STR_LIT>", dtype=dtype)<EOL>with tf.control_dependencies([<EOL>assert_util.assert_positive(<EOL>concentration, message="<STR_LIT>"),<EOL>assert_util.assert_positive(<EOL>scale, message="<STR_LIT>"),<EOL>] if validate_args else []):<EOL><INDENT>self._concentration = tf.identity(concentration, name="<STR_LIT>")<EOL>self._scale = tf.identity(scale, name="<STR_LIT>")<EOL><DEDENT>dtype_util.assert_same_float_dtype([self._concentration, self._scale])<EOL><DEDENT>super(InverseGamma, self).__init__(<EOL>dtype=self._concentration.dtype,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>parameters=parameters,<EOL>graph_parents=[self._concentration, self._scale],<EOL>name=name)<EOL>

Construct InverseGamma with `concentration` and `scale` parameters.

        The parameters `concentration` and `scale` must be shaped in a way that
        supports broadcasting (e.g. `concentration + scale` is a valid operation).

        Args:
          concentration: Floating point tensor, the concentration params of the
            distribution(s). Must contain only positive values.
          scale: Floating point tensor, the scale params of the distribution(s).
            Must contain only positive values.
          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.
          rate: Deprecated (mis-named) alias for `scale`.
          name: Python `str` name prefixed to Ops created by this class.

        Raises:
          TypeError: if `concentration` and `scale` are different dtypes.


return self._concentration<EOL>


return self._scale<EOL>


@kullback_leibler.RegisterKL(Chi, Chi)<EOL>def _kl_chi_chi(a, b, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>return (<NUM_LIT:0.5> * tf.math.digamma(<NUM_LIT:0.5> * a.df) * (a.df - b.df) +<EOL>tf.math.lgamma(<NUM_LIT:0.5> * b.df) - tf.math.lgamma(<NUM_LIT:0.5> * a.df))<EOL><DEDENT>

Calculate the batched KL divergence KL(a || b) with a and b Chi.

    Args:
      a: instance of a Chi distribution object.
      b: instance of a Chi distribution object.
      name: (optional) Name to use for created operations.
        default is "kl_chi_chi".

    Returns:
      Batchwise KL(a || b)


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>df = tf.convert_to_tensor(<EOL>value=df,<EOL>name="<STR_LIT>",<EOL>dtype=dtype_util.common_dtype([df], preferred_dtype=tf.float32))<EOL>validation_assertions = (<EOL>[assert_util.assert_positive(df)] if validate_args else [])<EOL>with tf.control_dependencies(validation_assertions):<EOL><INDENT>self._df = tf.identity(df, name="<STR_LIT>")<EOL><DEDENT>super(Chi, self).__init__(<EOL>distribution=chi2.Chi2(df=self._df,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats),<EOL>bijector=invert_bijector.Invert(square_bijector.Square()),<EOL>parameters=parameters,<EOL>name=name)<EOL><DEDENT>

Construct Chi distributions with parameter `df`.

        Args:
          df: Floating point tensor, the degrees of freedom of the
            distribution(s). `df` must contain only positive values.
          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: Python `str` name prefixed to Ops created by this class.
            Default value: `'Chi'`.


return self._df<EOL>

Distribution parameter for degrees of freedom.


output = tensor_to_broadcast<EOL>for tensor in target_tensors:<EOL><INDENT>output += tf.zeros_like(tensor)<EOL><DEDENT>return output<EOL>

Helper to broadcast a tensor using a list of target tensors.


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([low, high, peak], tf.float32)<EOL>low = tf.convert_to_tensor(value=low, name="<STR_LIT>", dtype=dtype)<EOL>high = tf.convert_to_tensor(value=high, name="<STR_LIT>", dtype=dtype)<EOL>peak = tf.convert_to_tensor(value=peak, name="<STR_LIT>", dtype=dtype)<EOL>with tf.control_dependencies([<EOL>assert_util.assert_less(<EOL>low, high, message="<STR_LIT>"),<EOL>assert_util.assert_less_equal(<EOL>low, peak, message="<STR_LIT>"),<EOL>assert_util.assert_less_equal(<EOL>peak, high, message="<STR_LIT>"),<EOL>] if validate_args else []):<EOL><INDENT>self._low = tf.identity(low, name="<STR_LIT>")<EOL>self._high = tf.identity(high, name="<STR_LIT>")<EOL>self._peak = tf.identity(peak, name="<STR_LIT>")<EOL>dtype_util.assert_same_float_dtype(<EOL>[self._low, self._high, self._peak])<EOL><DEDENT><DEDENT>super(Triangular, self).__init__(<EOL>dtype=self._low.dtype,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._low, self._high, self._peak],<EOL>name=name)<EOL>

Initialize a batch of Triangular distributions.

        Args:
          low: Floating point tensor, lower boundary of the output interval. Must
            have `low < high`.
            Default value: `0`.
          high: Floating point tensor, upper boundary of the output interval. Must
            have `low < high`.
            Default value: `1`.
          peak: Floating point tensor, mode of the output interval. Must have
            `low <= peak` and `peak <= high`.
            Default value: `0.5`.
          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.
            Default value: `False`.
          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.
            Default value: `True`.
          name: Python `str` name prefixed to Ops created by this class.
            Default value: `'Triangular'`.

        Raises:
          InvalidArgumentError: if `validate_args=True` and one of the following is
            True:
            * `low >= high`.
            * `peak > high`.
            * `low > peak`.


return self._low<EOL>


return self._high<EOL>


return self._peak<EOL>

Peak of the distribution. Lies in the interval.


return (self.peak - self.low) / (self.high - self.low)<EOL>


@kullback_leibler.RegisterKL(Beta, Beta)<EOL>def _kl_beta_beta(d1, d2, name=None):

def delta(fn, is_property=True):<EOL><INDENT>fn1 = getattr(d1, fn)<EOL>fn2 = getattr(d2, fn)<EOL>return (fn2 - fn1) if is_property else (fn2() - fn1())<EOL><DEDENT>with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>return (delta("<STR_LIT>", is_property=False) -<EOL>tf.math.digamma(d1.concentration1) * delta("<STR_LIT>") -<EOL>tf.math.digamma(d1.concentration0) * delta("<STR_LIT>") +<EOL>(tf.math.digamma(d1.total_concentration) *<EOL>delta("<STR_LIT>")))<EOL><DEDENT>

Calculate the batchwise KL divergence KL(d1 || d2) with d1 and d2 Beta.

    Args:
      d1: instance of a Beta distribution object.
      d2: instance of a Beta distribution object.
      name: (optional) Name to use for created operations.
        default is "kl_beta_beta".

    Returns:
      Batchwise KL(d1 || d2)


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([concentration1, concentration0],<EOL>tf.float32)<EOL>self._concentration1 = self._maybe_assert_valid_concentration(<EOL>tf.convert_to_tensor(<EOL>value=concentration1, name="<STR_LIT>", dtype=dtype),<EOL>validate_args)<EOL>self._concentration0 = self._maybe_assert_valid_concentration(<EOL>tf.convert_to_tensor(<EOL>value=concentration0, name="<STR_LIT>", dtype=dtype),<EOL>validate_args)<EOL>self._total_concentration = self._concentration1 + self._concentration0<EOL><DEDENT>super(Beta, self).__init__(<EOL>dtype=dtype,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>parameters=parameters,<EOL>graph_parents=[<EOL>self._concentration1, self._concentration0,<EOL>self._total_concentration<EOL>],<EOL>name=name)<EOL>

Initialize a batch of Beta distributions.

        Args:
          concentration1: Positive floating-point `Tensor` indicating mean
            number of successes; aka "alpha". Implies `self.dtype` and
            `self.batch_shape`, i.e.,
            `concentration1.shape = [N1, N2, ..., Nm] = self.batch_shape`.
          concentration0: Positive floating-point `Tensor` indicating mean
            number of failures; aka "beta". Otherwise has same semantics as
            `concentration1`.
          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: Python `str` name prefixed to Ops created by this class.


return self._concentration1<EOL>

Concentration parameter associated with a `1` outcome.


return self._concentration0<EOL>

Concentration parameter associated with a `0` outcome.


return self._total_concentration<EOL>


def _maybe_assert_valid_concentration(self, concentration, validate_args):

if not validate_args:<EOL><INDENT>return concentration<EOL><DEDENT>return distribution_util.with_dependencies([<EOL>assert_util.assert_positive(<EOL>concentration, message="<STR_LIT>"),<EOL>], concentration)<EOL>

Checks the validity of a concentration parameter.


if not self.validate_args:<EOL><INDENT>return x<EOL><DEDENT>return distribution_util.with_dependencies([<EOL>assert_util.assert_positive(x, message="<STR_LIT>"),<EOL>assert_util.assert_less(x, <NUM_LIT:1.>, message="<STR_LIT>"),<EOL>], x)<EOL>


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._logits, self._probs = distribution_util.get_logits_and_probs(<EOL>logits, probs, validate_args=validate_args, name=name)<EOL>with tf.control_dependencies(<EOL>[assert_util.assert_positive(self._probs)] if validate_args else []):<EOL><INDENT>self._probs = tf.identity(self._probs, name="<STR_LIT>")<EOL><DEDENT><DEDENT>super(Geometric, self).__init__(<EOL>dtype=self._probs.dtype,<EOL>reparameterization_type=reparameterization.NOT_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._probs, self._logits],<EOL>name=name)<EOL>

Construct Geometric distributions.

        Args:
          logits: Floating-point `Tensor` with shape `[B1, ..., Bb]` where `b >= 0`
            indicates the number of batch dimensions. Each entry represents logits
            for the probability of success for independent Geometric distributions
            and must be in the range `(-inf, inf]`. Only one of `logits` or `probs`
            should be specified.
          probs: Positive floating-point `Tensor` with shape `[B1, ..., Bb]`
            where `b >= 0` indicates the number of batch dimensions. Each entry
            represents the probability of success for independent Geometric
            distributions and must be in the range `(0, 1]`. Only one of `logits`
            or `probs` should be specified.
          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: Python `str` name prefixed to Ops created by this class.


return self._logits<EOL>


return self._probs<EOL>


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>self._distribution_fn = distribution_fn<EOL>self._sample0 = sample0<EOL>self._distribution0 = (distribution_fn() if sample0 is None<EOL>else distribution_fn(sample0))<EOL>if num_steps is None:<EOL><INDENT>num_steps = tensorshape_util.num_elements(<EOL>self._distribution0.event_shape)<EOL>if num_steps is None:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT><DEDENT>if num_steps < <NUM_LIT:1>:<EOL><INDENT>raise ValueError("<STR_LIT>".format(num_steps))<EOL><DEDENT>self._num_steps = num_steps<EOL><DEDENT>super(Autoregressive, self).__init__(<EOL>dtype=self._distribution0.dtype,<EOL>reparameterization_type=self._distribution0.reparameterization_type,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=self._distribution0._graph_parents,  <EOL>name=name)<EOL>

Construct an `Autoregressive` distribution.

        Args:
          distribution_fn: Python `callable` which constructs a
            `tfd.Distribution`-like instance from a `Tensor` (e.g.,
            `sample0`). The function must respect the "autoregressive property",
            i.e., there exists a permutation of event such that each coordinate is a
            diffeomorphic function of on preceding coordinates.
          sample0: Initial input to `distribution_fn`; used to
            build the distribution in `__init__` which in turn specifies this
            distribution's properties, e.g., `event_shape`, `batch_shape`, `dtype`.
            If unspecified, then `distribution_fn` should be default constructable.
          num_steps: Number of times `distribution_fn` is composed from samples,
            e.g., `num_steps=2` implies
            `distribution_fn(distribution_fn(sample0).sample(n)).sample()`.
          validate_args: Python `bool`.  Whether to validate input with asserts.
            If `validate_args` is `False`, and the inputs are invalid,
            correct behavior is not guaranteed.
          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: Python `str` name prefixed to Ops created by this class.
            Default value: "Autoregressive".

        Raises:
          ValueError: if `num_steps` and
            `num_elements(distribution_fn(sample0).event_shape)` are both `None`.
          ValueError: if `num_steps < 1`.


<EOL>eigenvalues, _ = tf.linalg.eigh(x)<EOL>return tf.cast(<EOL>tf.reduce_min(input_tensor=eigenvalues, axis=-<NUM_LIT:1>) >= <NUM_LIT:0>, dtype=x.dtype)<EOL>

Computes whether each square matrix in the input is positive semi-definite.

    Args:
      x: A floating-point `Tensor` of shape `[B1, ..., Bn, M, M]`.

    Returns:
      mask: A floating-point `Tensor` of shape `[B1, ... Bn]`.  Each
        scalar is 1 if the corresponding matrix was PSD, otherwise 0.


<EOL>return tf.cast(tf.linalg.det(x) > det_bounds, dtype=x.dtype)<EOL>

Returns whether the input matches the given determinant limit.

    Args:
      x: A floating-point `Tensor` of shape `[B1, ..., Bn, M, M]`.
      det_bounds: A floating-point `Tensor` that must broadcast to shape
        `[B1, ..., Bn]`, giving the desired lower bound on the
        determinants in `x`.

    Returns:
      mask: A floating-point `Tensor` of shape [B1, ..., Bn].  Each
        scalar is 1 if the corresponding matrix had determinant above
        the corresponding bound, otherwise 0.


def _uniform_correlation_like_matrix(num_rows, batch_shape, dtype, seed):

num_entries = num_rows * (num_rows + <NUM_LIT:1>) / <NUM_LIT:2><EOL>ones = tf.ones(shape=[num_entries], dtype=dtype)<EOL>unifs = uniform.Uniform(-ones, ones).sample(batch_shape, seed=seed)<EOL>tril = util.fill_triangular(unifs)<EOL>symmetric = tril + tf.linalg.matrix_transpose(tril)<EOL>diagonal_ones = tf.ones(<EOL>shape=util.pad(batch_shape, axis=<NUM_LIT:0>, back=True, value=num_rows),<EOL>dtype=dtype)<EOL>return tf.linalg.set_diag(symmetric, diagonal_ones)<EOL>

Returns a uniformly random `Tensor` of "correlation-like" matrices.

    A "correlation-like" matrix is a symmetric square matrix with all entries
    between -1 and 1 (inclusive) and 1s on the main diagonal.  Of these,
    the ones that are positive semi-definite are exactly the correlation
    matrices.

    Args:
      num_rows: Python `int` dimension of the correlation-like matrices.
      batch_shape: `Tensor` or Python `tuple` of `int` shape of the
        batch to return.
      dtype: `dtype` of the `Tensor` to return.
      seed: Random seed.

    Returns:
      matrices: A `Tensor` of shape `batch_shape + [num_rows, num_rows]`
        and dtype `dtype`.  Each entry is in [-1, 1], and each matrix
        along the bottom two dimensions is symmetric and has 1s on the
        main diagonal.


def correlation_matrix_volume_rejection_samples(<EOL>det_bounds, dim, sample_shape, dtype, seed):

with tf.compat.v1.name_scope("<STR_LIT>"):<EOL><INDENT>rej_proposals = _uniform_correlation_like_matrix(<EOL>dim, sample_shape, dtype, seed=seed)<EOL>rej_proposal_volume = <NUM_LIT> ** (dim * (dim - <NUM_LIT:1>) / <NUM_LIT>)<EOL>rej_weights = rej_proposal_volume * _psd_mask(<EOL>rej_proposals) * _det_large_enough_mask(rej_proposals, det_bounds)<EOL>return rej_weights, rej_proposal_volume<EOL><DEDENT>

Returns rejection samples from trying to get good correlation matrices.

    The proposal being rejected from is the uniform distribution on
    "correlation-like" matrices.  We say a matrix is "correlation-like"
    if it is a symmetric square matrix with all entries between -1 and 1
    (inclusive) and 1s on the main diagonal.  Of these, the ones that
    are positive semi-definite are exactly the correlation matrices.

    The rejection algorithm, then, is to sample a `Tensor` of
    `sample_shape` correlation-like matrices of dimensions `dim` by
    `dim`, and check each one for (i) being a correlation matrix (i.e.,
    PSD), and (ii) having determinant at least the corresponding entry
    of `det_bounds`.

    Args:
      det_bounds: A `Tensor` of lower bounds on the determinants of
        acceptable matrices.  The shape must broadcast with `sample_shape`.
      dim: A Python `int` dimension of correlation matrices to sample.
      sample_shape: Python `tuple` of `int` shape of the samples to
        compute, excluding the two matrix dimensions.
      dtype: The `dtype` in which to do the computation.
      seed: Random seed.

    Returns:
      weights: A `Tensor` of shape `sample_shape`.  Each entry is 0 if the
        corresponding matrix was not a correlation matrix, or had too
        small of a determinant.  Otherwise, the entry is the
        multiplicative inverse of the density of proposing that matrix
        uniformly, i.e., the volume of the set of `dim` by `dim`
        correlation-like matrices.
      volume: The volume of the set of `dim` by `dim` correlation-like
        matrices.


<EOL>if optimize is None or stats is None:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>")<EOL><DEDENT>if len(samples.shape) != <NUM_LIT:1>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>n = len(samples)<EOL>low = np.amin(samples)<EOL>high = np.amax(samples)<EOL>successes = np.count_nonzero(samples - low)<EOL>failures = np.count_nonzero(samples - high)<EOL>if successes + failures != n:<EOL><INDENT>uniques = np.unique(samples)<EOL>msg = ("<STR_LIT>"<EOL>"<STR_LIT>".format(uniques[<NUM_LIT:0>], uniques[<NUM_LIT:1>], uniques[<NUM_LIT:2>]))<EOL>raise ValueError(msg)<EOL><DEDENT>def p_small_enough(p):<EOL><INDENT>prob = stats.binom.logcdf(successes, n, p)<EOL>return prob - np.log(error_rate / <NUM_LIT>)<EOL><DEDENT>def p_big_enough(p):<EOL><INDENT>prob = stats.binom.logsf(successes, n, p)<EOL>return prob - np.log(error_rate / <NUM_LIT>)<EOL><DEDENT>high_p = optimize.brentq(<EOL>p_small_enough, float(successes) / n, <NUM_LIT:1.>, rtol=<NUM_LIT>)<EOL>low_p = optimize.brentq(<EOL>p_big_enough, <NUM_LIT:0.>, float(successes) / n, rtol=<NUM_LIT>)<EOL>low_interval = low + (high - low) * low_p<EOL>high_interval = low + (high - low) * high_p<EOL>return (low_interval, high_interval)<EOL>

Computes a confidence interval for the mean of the given 1-D distribution.

    Assumes (and checks) that the given distribution is Bernoulli, i.e.,
    takes only two values.  This licenses using the CDF of the binomial
    distribution for the confidence, which is tighter (for extreme
    probabilities) than the DKWM inequality.  The method is known as the
    [Clopper-Pearson method]
    (https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval).

    Assumes:

    - The given samples were drawn iid from the distribution of interest.

    - The given distribution is a Bernoulli, i.e., supported only on
      low and high.

    Guarantees:

    - The probability (over the randomness of drawing the given sample)
      that the true mean is outside the returned interval is no more
      than the given error_rate.

    Args:
      samples: `np.ndarray` of samples drawn iid from the distribution
        of interest.
      error_rate: Python `float` admissible rate of mistakes.

    Returns:
      low: Lower bound of confidence interval.
      high: Upper bound of confidence interval.

    Raises:
      ValueError: If `samples` has rank other than 1 (batch semantics
        are not implemented), or if `samples` contains values other than
        `low` or `high` (as that makes the distribution not Bernoulli).


def compute_true_volumes(<EOL>det_bounds, dim, num_samples, error_rate=<NUM_LIT>, seed=<NUM_LIT>):

bounds = {}<EOL>with tf.compat.v1.Session() as sess:<EOL><INDENT>rej_weights, _ = correlation_matrix_volume_rejection_samples(<EOL>det_bounds, dim, [num_samples, len(det_bounds)], np.float32, seed=seed)<EOL>rej_weights = sess.run(rej_weights)<EOL>for rw, det in zip(np.rollaxis(rej_weights, <NUM_LIT:1>), det_bounds):<EOL><INDENT>template = ("<STR_LIT>"<EOL>"<STR_LIT>")<EOL>print(template.format(dim, dim, det))<EOL>sys.stdout.flush()<EOL>bounds[det] = _clopper_pearson_confidence_interval(<EOL>rw, error_rate=error_rate)<EOL><DEDENT>return bounds<EOL><DEDENT>

Returns confidence intervals for the desired correlation matrix volumes.

    The confidence intervals are computed by the [Clopper-Pearson method]
    (https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval).

    Args:
      det_bounds: A rank-1 numpy array of lower bounds on the
        determinants of acceptable matrices.  Entries must be unique.
      dim: A Python `int` dimension of correlation matrices to sample.
      num_samples: The number of samples to draw.
      error_rate: The statistical significance of the returned
        confidence intervals.  The significance is broadcast: Each
        returned interval separately may be incorrect with probability
        (under the sample of correlation-like matrices drawn internally)
        at most `error_rate`.
      seed: Random seed.

    Returns:
      bounds: A Python `dict` mapping each determinant bound to the low, high
        tuple giving the confidence interval.


def _slice_single_param(param, param_event_ndims, slices, dist_batch_shape):

<EOL>param_shape = tf.shape(input=param)<EOL>insert_ones = tf.ones(<EOL>[tf.size(input=dist_batch_shape) + param_event_ndims - tf.rank(param)],<EOL>dtype=param_shape.dtype)<EOL>new_param_shape = tf.concat([insert_ones, param_shape], axis=<NUM_LIT:0>)<EOL>full_batch_param = tf.reshape(param, new_param_shape)<EOL>param_slices = []<EOL>param_dim_idx = <NUM_LIT:0><EOL>batch_dim_idx = <NUM_LIT:0><EOL>for slc in slices:<EOL><INDENT>if slc is tf.newaxis:<EOL><INDENT>param_slices.append(slc)<EOL>continue<EOL><DEDENT>if slc is Ellipsis:<EOL><INDENT>if batch_dim_idx < <NUM_LIT:0>:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(slices))<EOL><DEDENT>param_slices.append(slc)<EOL>num_remaining_non_newaxis_slices = sum(<EOL>[s is not tf.newaxis for s in slices[slices.index(Ellipsis) + <NUM_LIT:1>:]])<EOL>batch_dim_idx = -num_remaining_non_newaxis_slices<EOL>param_dim_idx = batch_dim_idx - param_event_ndims<EOL>continue<EOL><DEDENT>param_dim_size = new_param_shape[param_dim_idx]<EOL>batch_dim_size = dist_batch_shape[batch_dim_idx]<EOL>is_broadcast = batch_dim_size > param_dim_size<EOL>if isinstance(slc, slice):<EOL><INDENT>start, stop, step = slc.start, slc.stop, slc.step<EOL>if start is not None:<EOL><INDENT>start = tf.where(is_broadcast, <NUM_LIT:0>, start)<EOL><DEDENT>if stop is not None:<EOL><INDENT>stop = tf.where(is_broadcast, <NUM_LIT:1>, stop)<EOL><DEDENT>if step is not None:<EOL><INDENT>step = tf.where(is_broadcast, <NUM_LIT:1>, step)<EOL><DEDENT>param_slices.append(slice(start, stop, step))<EOL><DEDENT>else:  <EOL><INDENT>param_slices.append(tf.where(is_broadcast, <NUM_LIT:0>, slc))<EOL><DEDENT>param_dim_idx += <NUM_LIT:1><EOL>batch_dim_idx += <NUM_LIT:1><EOL><DEDENT>param_slices.extend([ALL_SLICE] * param_event_ndims)<EOL>return full_batch_param.__getitem__(param_slices)<EOL>

Slices a single parameter of a distribution.

    Args:
      param: A `Tensor`, the original parameter to slice.
      param_event_ndims: `int` event parameterization rank for this parameter.
      slices: A `tuple` of normalized slices.
      dist_batch_shape: The distribution's batch shape `Tensor`.

    Returns:
      new_param: A `Tensor`, batch-sliced according to slices.


def _slice_params_to_dict(dist, params_event_ndims, slices):

override_dict = {}<EOL>for param_name, param_event_ndims in six.iteritems(params_event_ndims):<EOL><INDENT>if param_name not in dist.parameters:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'.format(dist, param_name))<EOL><DEDENT>param = dist.parameters[param_name]<EOL>if param is None:<EOL><INDENT>continue<EOL><DEDENT>dtype = None<EOL>if hasattr(dist, param_name):<EOL><INDENT>attr = getattr(dist, param_name)<EOL>dtype = getattr(attr, '<STR_LIT>', None)<EOL><DEDENT>if dtype is None:<EOL><INDENT>dtype = dist.dtype<EOL>warnings.warn('<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>param_name, dist, dtype))<EOL><DEDENT>param = tf.convert_to_tensor(value=param, dtype=dtype)<EOL>override_dict[param_name] = _slice_single_param(param, param_event_ndims,<EOL>slices,<EOL>dist.batch_shape_tensor())<EOL><DEDENT>return override_dict<EOL>

Computes the override dictionary of sliced parameters.

    Args:
      dist: The tfd.Distribution being batch-sliced.
      params_event_ndims: Per-event parameter ranks, a `str->int` `dict`.
      slices: Slices as received by __getitem__.

    Returns:
      overrides: `str->Tensor` `dict` of batch-sliced parameter overrides.


def _apply_single_step(dist, params_event_ndims, slices, params_overrides):

if len(slices) == <NUM_LIT:1> and slices[<NUM_LIT:0>] == Ellipsis:<EOL><INDENT>override_dict = {}<EOL><DEDENT>else:<EOL><INDENT>override_dict = _slice_params_to_dict(dist, params_event_ndims, slices)<EOL><DEDENT>override_dict.update(params_overrides)<EOL>parameters = dict(dist.parameters, **override_dict)<EOL>new_dist = type(dist)(**parameters)<EOL>return new_dist<EOL>

Applies a single slicing step to `dist`, returning a new instance.


def _apply_slice_sequence(dist, params_event_ndims, slice_overrides_seq):

for slices, overrides in slice_overrides_seq:<EOL><INDENT>dist = _apply_single_step(dist, params_event_ndims, slices, overrides)<EOL><DEDENT>return dist<EOL>

Applies a sequence of slice or copy-with-overrides operations to `dist`.


def batch_slice(dist, params_event_ndims, params_overrides, slices):

if not isinstance(slices, collections.Sequence):<EOL><INDENT>slices = (slices,)<EOL><DEDENT>orig_dist, slice_overrides_seq = getattr(dist, PROVENANCE_ATTR, (dist, []))<EOL>slice_overrides_seq += [(slices, params_overrides)]<EOL>dist = _apply_slice_sequence(orig_dist, params_event_ndims,<EOL>slice_overrides_seq)<EOL>setattr(dist, PROVENANCE_ATTR, (orig_dist, slice_overrides_seq))<EOL>return dist<EOL>

Slices `dist` along its batch dimensions. Helper for tfd.Distribution.

    Args:
      dist: A `tfd.Distribution` instance.
      params_event_ndims: A `dict` of `str->int` indicating the number of
        dimensions of a given parameter required to parameterize a single event.
      params_overrides: A `dict` of parameter overrides. (e.g. from
        `Distribution.copy`).
      slices: A `slice` or `int` or `int` `Tensor` or `tf.newaxis` or `tuple`
        thereof. (e.g. the argument of a `__getitem__` method).

    Returns:
      new_dist: A batch-sliced `tfd.Distribution`.


def assert_true_cdf_equal_by_dkwm(<EOL>samples, cdf, left_continuous_cdf=None, false_fail_rate=<NUM_LIT>, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[samples, false_fail_rate]):<EOL><INDENT>dtype = dtype_util.common_dtype([samples, false_fail_rate], tf.float32)<EOL>samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>', dtype=dtype)<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>tf.compat.v1.assert_scalar(false_fail_rate)  <EOL>itemwise_false_fail_rate = _itemwise_error_rate(<EOL>total_rate=false_fail_rate,<EOL>param_tensors=[], samples_tensor=samples)<EOL>n = tf.shape(input=samples)[<NUM_LIT:0>]<EOL>envelope = _dkwm_cdf_envelope(n, itemwise_false_fail_rate)<EOL>distance = kolmogorov_smirnov_distance(samples, cdf, left_continuous_cdf)<EOL>return tf.compat.v1.assert_less_equal(<EOL>distance, envelope, message='<STR_LIT>')<EOL><DEDENT>

Asserts the full CDF of the given distribution is as expected.

    More precisely, fails if there is enough evidence (using the
    [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval))
    that the true CDF of some distribution from which the given samples are
    drawn is _not_ the given expected CDF with statistical significance
    `false_fail_rate` or stronger, otherwise passes.  If you also want to
    check that you are gathering enough evidence that a pass is not
    spurious, see `min_num_samples_for_dkwm_cdf_test` and
    `min_discrepancy_of_true_cdfs_detectable_by_dkwm`.

    If the distribution in question has atoms (e.g., is discrete), computing this
    test requires CDF values for both sides of the discontinuity.  In this case,
    the `cdf` argument is assumed to compute the CDF inclusive of the atom, i.e.,
    cdf(x) = Pr(X <= x).  The user must also supply the `left_continuous_cdf`,
    which must compute the cdf exclusive of the atom, i.e., left_continuous_cdf(x)
    = Pr(X < x).  Invariant: cdf(x) - left_continuous_cdf(x) = pmf(x).

    For example, the two required cdfs of the degenerate distribution that places
    all the mass at 0 can be given as
    ```
    cdf=lambda x: tf.where(x < 0, 0., 1.)
    left_continuous_cdf=lambda x: tf.where(x <= 0, 0., 1.)
    ```

    Note that `false_fail_rate` is a total false failure rate for all
    the assertions in the batch.  As such, if the batch is nontrivial,
    the assertion will insist on stronger evidence to fail any one member.

    Args:
      samples: Tensor of shape [n] + B.  Samples from some (batch of) scalar-event
        distribution(s) of interest, giving a (batch of) empirical CDF(s).
        Assumed IID across the 0 dimension.
      cdf: Analytic cdf inclusive of any atoms, as a function that can compute CDF
        values in batch.  Must accept a Tensor of shape B + [n] and the same dtype
        as `samples` and return a Tensor of shape B + [n] of CDF values.  For each
        sample x, `cdf(x) = Pr(X <= x)`.
      left_continuous_cdf: Analytic left-continuous cdf, as a function that can
        compute CDF values in batch.  Must accept a Tensor of shape B + [n] and
        the same dtype as `samples` and return a Tensor of shape B + [n] of CDF
        values.  For each sample x, `left_continuous_cdf(x) = Pr(X < x)`.  If the
        distribution under test has no atoms (i.e., the CDF is continuous), this
        is redundant and may be omitted.  Conversely, if this argument is omitted,
        the test assumes the distribution is atom-free.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of mistakes.
      name: A name for this operation (optional).

    Returns:
      check: Op that raises `InvalidArgumentError` if any expected CDF is
        outside the corresponding confidence envelope.


def min_discrepancy_of_true_cdfs_detectable_by_dkwm(<EOL>n, false_fail_rate, false_pass_rate, name=None):

with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[n, false_fail_rate, false_pass_rate]):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[n, false_fail_rate, false_pass_rate], tf.float32)<EOL>n = tf.convert_to_tensor(value=n, name='<STR_LIT:n>', dtype=dtype)<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>false_pass_rate = tf.convert_to_tensor(<EOL>value=false_pass_rate, name='<STR_LIT>', dtype=dtype)<EOL>sampling_envelope = _dkwm_cdf_envelope(n, false_pass_rate)<EOL>itemwise_false_fail_rate = _itemwise_error_rate(<EOL>total_rate=false_fail_rate, param_tensors=[n])<EOL>analysis_envelope = _dkwm_cdf_envelope(n, itemwise_false_fail_rate)<EOL>return sampling_envelope + analysis_envelope<EOL><DEDENT>

Returns the minimum CDF discrepancy that a DKWM-based test can detect.

    DKWM is the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval).

    Note that `false_fail_rate` is a total false failure rate for all
    the tests in the batch.  As such, if the batch is nontrivial, each
    member will demand more samples.  The `false_pass_rate` is also
    interpreted as a total, but is treated asymmetrically: If each test
    in the batch detects its corresponding discrepancy with probability
    at least `1 - false_pass_rate`, then running all those tests and
    failing if any one fails will jointly detect all those discrepancies
    with the same `false_pass_rate`.

    Args:
      n: `Tensor` of numbers of samples to be drawn from the distributions
        of interest.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of false failures.
      false_pass_rate: *Scalar* floating-point `Tensor` admissible rate
        of false passes.
      name: A name for this operation (optional).

    Returns:
      discr: `Tensor` of lower bounds on the K-S distances between true
         CDFs detectable by a DKWM-based test.

    For each batch member `i`, of `K` total, drawing `n[i]` samples from some
    scalar distribution is enough to detect a K-S distance in CDFs of size
    `discr[i]` or more.  Specifically, we guarantee that (a) if the true CDF is
    the expected CDF, then `assert_true_cdf_equal_by_dkwm` will fail with
    probability at most `false_fail_rate / K` (which amounts to `false_fail_rate`
    if applied to the whole batch at once), and (b) if the true CDF differs from
    the expected CDF by at least `discr[i]`, `assert_true_cdf_equal_by_dkwm` will
    pass with probability at most `false_pass_rate`.

    The detectable discrepancy scales as

    - `O(1 / sqrt(n[i]))`,
    - `O(-log(false_fail_rate/K))`, and
    - `O(-log(false_pass_rate))`.


def kolmogorov_smirnov_distance(<EOL>samples, cdf, left_continuous_cdf=None, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [samples]):<EOL><INDENT>dtype = dtype_util.common_dtype([samples], tf.float32)<EOL>samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>', dtype=dtype)<EOL>samples = _move_dim_and_sort(samples)<EOL>cdfs = cdf(samples)<EOL>if left_continuous_cdf is None:<EOL><INDENT>left_continuous_cdfs = cdfs<EOL><DEDENT>else:<EOL><INDENT>left_continuous_cdfs = left_continuous_cdf(samples)<EOL><DEDENT>n = tf.cast(tf.shape(input=samples)[-<NUM_LIT:1>], dtype=cdfs.dtype)<EOL>low_empirical_cdfs = tf.range(n, dtype=cdfs.dtype) / n<EOL>high_empirical_cdfs = tf.range(<NUM_LIT:1>, n+<NUM_LIT:1>, dtype=cdfs.dtype) / n<EOL>low_distances = tf.reduce_max(<EOL>input_tensor=left_continuous_cdfs - low_empirical_cdfs, axis=-<NUM_LIT:1>)<EOL>high_distances = tf.reduce_max(<EOL>input_tensor=high_empirical_cdfs - cdfs, axis=-<NUM_LIT:1>)<EOL>return tf.maximum(low_distances, high_distances)<EOL><DEDENT>

Computes the Kolmogorov-Smirnov distance between the given CDFs.

    The (absolute) Kolmogorov-Smirnov distance is the maximum (absolute)
    discrepancy between the CDFs, i.e.,

      sup_x(|cdf1(x) - cdf2(x)|)

    This is tractable to compute exactly when at least one CDF in question is an
    empirical CDF given by samples, because the analytic one need only be queried
    at the sampled values.

    If the distribution in question has atoms (e.g., is discrete), computing the
    distance requires CDF values for both sides of the discontinuity.  In this
    case, the `cdf` argument is assumed to compute the CDF inclusive of the atom,
    i.e., cdf(x) = Pr(X <= x).  The user must also supply the
    `left_continuous_cdf`, which must compute the cdf exclusive of the atom, i.e.,
    left_continuous_cdf(x) = Pr(X < x).

    For example, the two required cdfs of the degenerate distribution that places
    all the mass at 0 can be given as
    ```
    cdf=lambda x: tf.where(x < 0, 0., 1.)
    left_continuous_cdf=lambda x: tf.where(x <= 0, 0., 1.)
    ```

    Args:
      samples: Tensor of shape [n] + B.  Samples from some (batch of) scalar-event
        distribution(s) of interest, giving a (batch of) empirical CDF(s).
        Assumed IID across the 0 dimension.
      cdf: Analytic cdf inclusive of any atoms, as a function that can compute CDF
        values in batch.  Must accept a Tensor of shape B + [n] and the same dtype
        as `samples` and return a Tensor of shape B + [n] of CDF values.  For each
        sample x, `cdf(x) = Pr(X <= x)`.
      left_continuous_cdf: Analytic left-continuous cdf, as a function that can
        compute CDF values in batch.  Must accept a Tensor of shape B + [n] and
        the same dtype as `samples` and return a Tensor of shape B + [n] of CDF
        values.  For each sample x, `left_continuous_cdf(x) = Pr(X < x)`.  If the
        distribution under test has no atoms (i.e., the CDF is continuous), this
        is redundant and may be omitted.  Conversely, if this argument is omitted,
        the test assumes the distribution is atom-free.
      name: A name for this operation (optional).

    Returns:
      distance: Tensor of shape B: (Absolute) Kolmogorov-Smirnov distance between
        the empirical and analytic CDFs.


def kolmogorov_smirnov_distance_two_sample(samples1, samples2, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[samples1, samples2]):<EOL><INDENT>dtype = dtype_util.common_dtype([samples1, samples2], tf.float32)<EOL>samples1 = tf.convert_to_tensor(<EOL>value=samples1, name='<STR_LIT>', dtype=dtype)<EOL>samples2 = tf.convert_to_tensor(<EOL>value=samples2, name='<STR_LIT>', dtype=dtype)<EOL>samples2 = _move_dim_and_sort(samples2)<EOL>cdf = functools.partial(<EOL>empirical_cdfs, samples2,<EOL>continuity='<STR_LIT:right>', dtype=samples1.dtype)<EOL>left_continuous_cdf = functools.partial(<EOL>empirical_cdfs, samples2,<EOL>continuity='<STR_LIT:left>', dtype=samples1.dtype)<EOL>return kolmogorov_smirnov_distance(samples1, cdf, left_continuous_cdf)<EOL><DEDENT>

Computes the Kolmogorov-Smirnov distance between the given empirical CDFs.

    The (absolute) Kolmogorov-Smirnov distance is the maximum (absolute)
    discrepancy between the CDFs, i.e.,

      sup_x(|cdf1(x) - cdf2(x)|)

    This is tractable to compute exactly for empirical CDFs, because they are
    piecewise constant with known piece boundaries (the samples).

    This function works even if the samples have duplicates (e.g., if the
    underlying distribution is discrete).

    Args:
      samples1: Tensor of shape [n] + B.  Samples from some (batch of)
        scalar-event distribution(s) of interest, giving a (batch of) empirical
        CDF(s).  Assumed IID across the 0 dimension.
      samples2: Tensor of shape [m] + B.  Samples from some (batch of)
        scalar-event distribution(s) of interest, giving a (batch of) empirical
        CDF(s).  Assumed IID across the 0 dimension.
      name: A name for this operation (optional).

    Returns:
      distance: Tensor of shape B: (Absolute) Kolmogorov-Smirnov distance between
        the two empirical CDFs given by the samples.


<EOL>samples = distribution_util.move_dimension(samples, <NUM_LIT:0>, -<NUM_LIT:1>)<EOL>samples = _batch_sort_vector(samples)<EOL>return samples<EOL>

Internal helper for K-S distance computation.


def _batch_sort_vector(x, ascending=True, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [x]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>n = tf.shape(input=x)[-<NUM_LIT:1>]<EOL>if ascending:<EOL><INDENT>y, _ = tf.nn.top_k(-x, k=n, sorted=True)<EOL>y = -y<EOL><DEDENT>else:<EOL><INDENT>y, _ = tf.nn.top_k(x, k=n, sorted=True)<EOL><DEDENT>tensorshape_util.set_shape(y, x.shape)<EOL>return y<EOL><DEDENT>

Batch sort.  Sorts the -1 dimension of each batch member independently.


def empirical_cdfs(samples, positions, continuity='<STR_LIT:right>',<EOL>dtype=tf.float32, name=None):

if continuity not in ['<STR_LIT:left>', '<STR_LIT:right>']:<EOL><INDENT>msg = '<STR_LIT>'.format(<EOL>continuity)<EOL>raise ValueError(msg)<EOL><DEDENT>with tf.compat.v1.name_scope(name, '<STR_LIT>', [samples, positions]):<EOL><INDENT>n = tf.cast(tf.shape(input=samples)[-<NUM_LIT:1>], dtype=dtype)<EOL>indexes = tf.searchsorted(<EOL>sorted_sequence=samples, values=positions, side=continuity)<EOL>return tf.cast(indexes, dtype=dtype) / n<EOL><DEDENT>

Evaluates the empirical CDF of a batch of samples at a batch of positions.

    If elements of `positions` might be exactly equal to elements of `samples`
    (e.g., if the underlying distribution of interest is discrete), there is a
    difference between the conventional, right-continuous CDF (Pr[X <= x]) and a
    left-continuous variant (Pr[X < x]).  The latter can be accessed by setting
    `continuity='left'`.  The difference between the right-continuous and
    left-continuous CDFs is the empirical pmf at each point, i.e., how many times
    each element of `positions` occurs in its batch of `samples`.

    Note: Returns results parallel to `positions`, i.e., the values of the
    empirical CDF at those points.

    Note: The sample dimension is _last_, and the samples must be _sorted_ within
    each batch.

    Args:
      samples: Tensor of shape `batch + [num_samples]` of samples.  The samples
        must be in ascending order within each batch member.
      positions: Tensor of shape `batch + [m]` of positions where to evaluate the
        CDFs.  The positions need not be sorted.
      continuity: Whether to return a conventional, right-continuous CDF
        (`continuity = 'right'`, default) or a left-continuous CDF (`continuity =
        'left'`).  The value at each point `x` will be `F_n(X <= x)` or
        `F_n(X < x)`, respectively.
      dtype: dtype at which to evaluate the desired empirical CDFs.
      name: A name for this operation (optional).

    Returns:
      cdf: Tensor parallel to `positions`.  For each x in `positions`, gives the
        (right- or left-continuous, per the `continuity` argument) cdf at that
        position.  If `positions` contains duplicates, `cdf` will give each the
        same value.


def _do_maximum_mean(samples, envelope, high, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[samples, envelope, high]):<EOL><INDENT>dtype = dtype_util.common_dtype([samples, envelope, high], tf.float32)<EOL>samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>', dtype=dtype)<EOL>envelope = tf.convert_to_tensor(<EOL>value=envelope, name='<STR_LIT>', dtype=dtype)<EOL>high = tf.convert_to_tensor(value=high, name='<STR_LIT>', dtype=dtype)<EOL>n = tf.rank(samples)<EOL>perm = tf.concat([tf.range(<NUM_LIT:1>, n), [<NUM_LIT:0>]], axis=<NUM_LIT:0>)<EOL>samples = tf.transpose(a=samples, perm=perm)<EOL>samples = _batch_sort_vector(samples)<EOL>batch_size = tf.shape(input=samples)[-<NUM_LIT:1>]<EOL>batch_size = tf.cast(batch_size, dtype=dtype)<EOL>step = <NUM_LIT:1.> / batch_size<EOL>cum_steps = step * tf.range(<NUM_LIT:1>, batch_size + <NUM_LIT:1>, dtype=dtype)<EOL>max_mean_contrib = tf.clip_by_value(<EOL>cum_steps - envelope[..., tf.newaxis],<EOL>clip_value_min=<NUM_LIT:0.>,<EOL>clip_value_max=step)<EOL>return tf.reduce_sum(<EOL>input_tensor=samples * max_mean_contrib, axis=-<NUM_LIT:1>) + envelope * high<EOL><DEDENT>

Common code between maximum_mean and minimum_mean.


def assert_true_cdf_equal_by_dkwm_two_sample(<EOL>samples1, samples2, false_fail_rate=<NUM_LIT>, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[samples1, samples2, false_fail_rate]):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[samples1, samples2, false_fail_rate], tf.float32)<EOL>samples1 = tf.convert_to_tensor(<EOL>value=samples1, name='<STR_LIT>', dtype=dtype)<EOL>samples2 = tf.convert_to_tensor(<EOL>value=samples2, name='<STR_LIT>', dtype=dtype)<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>tf.compat.v1.assert_scalar(false_fail_rate)  <EOL>compatible_samples = tf.compat.v1.assert_equal(<EOL>tf.shape(input=samples1)[<NUM_LIT:1>:],<EOL>tf.shape(input=samples2)[<NUM_LIT:1>:])<EOL>with tf.control_dependencies([compatible_samples]):<EOL><INDENT>itemwise_false_fail_rate = _itemwise_error_rate(<EOL>total_rate=false_fail_rate,<EOL>param_tensors=[], samples_tensor=samples1)<EOL>n1 = tf.shape(input=samples1)[<NUM_LIT:0>]<EOL>envelope1 = _dkwm_cdf_envelope(n1, itemwise_false_fail_rate)<EOL>n2 = tf.shape(input=samples2)[<NUM_LIT:0>]<EOL>envelope2 = _dkwm_cdf_envelope(n2, itemwise_false_fail_rate)<EOL>distance = kolmogorov_smirnov_distance_two_sample(samples1, samples2)<EOL>return tf.compat.v1.assert_less_equal(<EOL>distance, envelope1 + envelope2,<EOL>message='<STR_LIT>')<EOL><DEDENT><DEDENT>

Asserts the full CDFs of the two given distributions are equal.

    More precisely, fails if there is enough evidence (using the
    [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval))
    that the true CDF of the distribution from which `samples1` are drawn is _not_
    the true CDF of the distribution from which `samples2` are drawn, with
    statistical significance `false_fail_rate` or stronger, otherwise passes.  If
    you also want to check that you are gathering enough evidence that a pass is
    not spurious, see `min_num_samples_for_dkwm_cdf_two_sample_test` and
    `min_discrepancy_of_true_cdfs_detectable_by_dkwm_two_sample`.

    This test works as written even if the distribution in question has atoms
    (e.g., is discrete).

    Note that `false_fail_rate` is a total false failure rate for all
    the assertions in the batch.  As such, if the batch is nontrivial,
    the assertion will insist on stronger evidence to fail any one member.

    Args:
      samples1: Tensor of shape [n] + B.  Samples from some (batch of)
        scalar-event distribution(s) of interest, giving a (batch of) empirical
        CDF(s).  Assumed IID across the 0 dimension.
      samples2: Tensor of shape [m] + B.  Samples from some (batch of)
        scalar-event distribution(s) of interest, giving a (batch of) empirical
        CDF(s).  Assumed IID across the 0 dimension.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of mistakes.
      name: A name for this operation (optional).

    Returns:
      check: Op that raises `InvalidArgumentError` if any expected CDF is
        outside the corresponding confidence envelope.


def min_discrepancy_of_true_cdfs_detectable_by_dkwm_two_sample(<EOL>n1, n2, false_fail_rate, false_pass_rate, name=None):

with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[n1, n2, false_fail_rate, false_pass_rate]):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[n1, n2, false_fail_rate, false_pass_rate], tf.float32)<EOL>n1 = tf.convert_to_tensor(value=n1, name='<STR_LIT>', dtype=dtype)<EOL>n2 = tf.convert_to_tensor(value=n2, name='<STR_LIT>', dtype=dtype)<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>false_pass_rate = tf.convert_to_tensor(<EOL>value=false_pass_rate, name='<STR_LIT>', dtype=dtype)<EOL>d1 = min_discrepancy_of_true_cdfs_detectable_by_dkwm(<EOL>n1, false_fail_rate / <NUM_LIT>, false_pass_rate / <NUM_LIT>)<EOL>d2 = min_discrepancy_of_true_cdfs_detectable_by_dkwm(<EOL>n2, false_fail_rate / <NUM_LIT>, false_pass_rate / <NUM_LIT>)<EOL>return d1 + d2<EOL><DEDENT>

Returns the minimum CDF discrepancy that a two-sample DKWM test can detect.

    DKWM is the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval).

    Note that `false_fail_rate` is a total false failure rate for all
    the tests in the batch.  As such, if the batch is nontrivial, each
    member will demand more samples.  The `false_pass_rate` is also
    interpreted as a total, but is treated asymmetrically: If each test
    in the batch detects its corresponding discrepancy with probability
    at least `1 - false_pass_rate`, then running all those tests and
    failing if any one fails will jointly detect all those discrepancies
    with the same `false_pass_rate`.

    Args:
      n1: `Tensor` of numbers of samples to be drawn from the distributions A.
      n2: `Tensor` of numbers of samples to be drawn from the distributions B.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of false failures.
      false_pass_rate: *Scalar* floating-point `Tensor` admissible rate
        of false passes.
      name: A name for this operation (optional).

    Returns:
      discr: `Tensor` of lower bounds on the K-S distances between true
         CDFs detectable by a DKWM-based test.

    For each batch member `i`, of `K` total, drawing `n1[i]` samples from scalar
    distribution A and `n2[i]` samples from scalar distribution B is enough to
    detect a K-S distance in CDFs of size `discr[i]` or more.  Specifically, we
    guarantee that (a) if their true CDFs are the same, then
    `assert_true_cdf_equal_by_dkwm_two_sample` will fail with probability at most
    `false_fail_rate / K` (which amounts to `false_fail_rate` if applied to the
    whole batch at once), and (b) if their true CDFs differ by at least
    `discr[i]`, `assert_true_cdf_equal_by_dkwm_two_sample` will pass with
    probability at most `false_pass_rate`.

    The detectable discrepancy scales as

    - `O(1 / sqrt(n[i]))`,
    - `O(-log(false_fail_rate/K))`, and
    - `O(-log(false_pass_rate))`.


def _maximum_mean(samples, envelope, high, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [samples, envelope, high]):<EOL><INDENT>dtype = dtype_util.common_dtype([samples, envelope, high], tf.float32)<EOL>samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>', dtype=dtype)<EOL>envelope = tf.convert_to_tensor(<EOL>value=envelope, name='<STR_LIT>', dtype=dtype)<EOL>high = tf.convert_to_tensor(value=high, name='<STR_LIT>', dtype=dtype)<EOL>xmax = tf.reduce_max(input_tensor=samples, axis=[<NUM_LIT:0>])<EOL>msg = '<STR_LIT>'<EOL>check_op = tf.compat.v1.assert_less_equal(xmax, high, message=msg)<EOL>with tf.control_dependencies([check_op]):<EOL><INDENT>return tf.identity(_do_maximum_mean(samples, envelope, high))<EOL><DEDENT><DEDENT>

Returns a stochastic upper bound on the mean of a scalar distribution.

    The idea is that if the true CDF is within an `eps`-envelope of the
    empirical CDF of the samples, and the support is bounded above, then
    the mean is bounded above as well.  In symbols,

    ```none
    sup_x(|F_n(x) - F(x)|) < eps
    ```

    The 0th dimension of `samples` is interpreted as independent and
    identically distributed samples.  The remaining dimensions are
    broadcast together with `envelope` and `high`, and operated on
    separately.

    Args:
      samples: Floating-point `Tensor` of samples from the distribution(s)
        of interest.  Entries are assumed IID across the 0th dimension.
        The other dimensions must broadcast with `envelope` and `high`.
      envelope: Floating-point `Tensor` of sizes of admissible CDF
        envelopes (i.e., the `eps` above).
      high: Floating-point `Tensor` of upper bounds on the distributions'
        supports.  `samples <= high`.
      name: A name for this operation (optional).

    Returns:
      bound: Floating-point `Tensor` of upper bounds on the true means.

    Raises:
      InvalidArgumentError: If some `sample` is found to be larger than
        the corresponding `high`.


def _minimum_mean(samples, envelope, low, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [samples, envelope, low]):<EOL><INDENT>dtype = dtype_util.common_dtype([samples, envelope, low], tf.float32)<EOL>samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>', dtype=dtype)<EOL>envelope = tf.convert_to_tensor(<EOL>value=envelope, name='<STR_LIT>', dtype=dtype)<EOL>low = tf.convert_to_tensor(value=low, name='<STR_LIT>', dtype=dtype)<EOL>xmin = tf.reduce_min(input_tensor=samples, axis=[<NUM_LIT:0>])<EOL>msg = '<STR_LIT>'<EOL>check_op = tf.compat.v1.assert_greater_equal(xmin, low, message=msg)<EOL>with tf.control_dependencies([check_op]):<EOL><INDENT>return - _do_maximum_mean(-samples, envelope, -low)<EOL><DEDENT><DEDENT>

Returns a stochastic lower bound on the mean of a scalar distribution.

    The idea is that if the true CDF is within an `eps`-envelope of the
    empirical CDF of the samples, and the support is bounded below, then
    the mean is bounded below as well.  In symbols,

    ```none
    sup_x(|F_n(x) - F(x)|) < eps
    ```

    The 0th dimension of `samples` is interpreted as independent and
    identically distributed samples.  The remaining dimensions are
    broadcast together with `envelope` and `low`, and operated on
    separately.

    Args:
      samples: Floating-point `Tensor` of samples from the distribution(s)
        of interest.  Entries are assumed IID across the 0th dimension.
        The other dimensions must broadcast with `envelope` and `low`.
      envelope: Floating-point `Tensor` of sizes of admissible CDF
        envelopes (i.e., the `eps` above).
      low: Floating-point `Tensor` of lower bounds on the distributions'
        supports.  `samples >= low`.
      name: A name for this operation (optional).

    Returns:
      bound: Floating-point `Tensor` of lower bounds on the true means.

    Raises:
      InvalidArgumentError: If some `sample` is found to be smaller than
        the corresponding `low`.


def _dkwm_cdf_envelope(n, error_rate, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [n, error_rate]):<EOL><INDENT>n = tf.cast(n, dtype=error_rate.dtype)<EOL>return tf.sqrt(-tf.math.log(error_rate / <NUM_LIT>) / (<NUM_LIT> * n))<EOL><DEDENT>

Computes the CDF envelope that the DKWM inequality licenses.

    The [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval)
    gives a stochastic bound on the distance between the true cumulative
    distribution function (CDF) of any distribution and its empirical
    CDF.  To wit, for `n` iid samples from any distribution with CDF F,

    ```none
    P(sup_x |F_n(x) - F(x)| > eps) < 2exp(-2n eps^2)
    ```

    This function computes the envelope size `eps` as a function of the
    number of samples `n` and the desired limit on the left-hand
    probability above.

    Args:
      n: `Tensor` of numbers of samples drawn.
      error_rate: Floating-point `Tensor` of admissible rates of mistakes.
      name: A name for this operation (optional).

    Returns:
      eps: `Tensor` of maximum distances the true CDF can be from the
        empirical CDF.  This scales as `O(sqrt(-log(error_rate)))` and
        as `O(1 / sqrt(n))`.  The shape is the broadcast of `n` and
        `error_rate`.


def check(t):<EOL><INDENT>samples_batch_shape = tf.shape(input=samples)[<NUM_LIT:1>:]<EOL>broadcasted_batch_shape = tf.broadcast_dynamic_shape(<EOL>samples_batch_shape, tf.shape(input=t))<EOL>samples_batch_ndims = tf.size(input=samples_batch_shape)<EOL>ge = tf.compat.v1.assert_greater_equal(samples_batch_ndims, tf.rank(t))<EOL>eq = tf.compat.v1.assert_equal(samples_batch_shape, broadcasted_batch_shape)<EOL>return ge, eq<EOL><DEDENT>checks = list(itertools.chain(*[check(t) for t in parameters]))<EOL>with tf.control_dependencies(checks):<EOL><INDENT>return tf.identity(samples)<EOL><DEDENT>

Check that broadcasting `samples` against `parameters` does not expand it.

    Why?  To be very sure that the samples tensor is not accidentally enlarged by
    broadcasting against tensors that are supposed to be describing the
    distribution(s) sampled from, lest the sample counts end up inflated.

    Args:
      samples: A `Tensor` whose shape is to be protected against broadcasting.
      parameters: A list of `Tensor`s who are parameters for the statistical test.

    Returns:
      samples: Return original `samples` with control dependencies attached
        to ensure no broadcasting.


def true_mean_confidence_interval_by_dkwm(<EOL>samples, low, high, error_rate=<NUM_LIT>, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[samples, low, high, error_rate]):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[samples, low, high, error_rate], tf.float32)<EOL>samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>', dtype=dtype)<EOL>low = tf.convert_to_tensor(value=low, name='<STR_LIT>', dtype=dtype)<EOL>high = tf.convert_to_tensor(value=high, name='<STR_LIT>', dtype=dtype)<EOL>error_rate = tf.convert_to_tensor(<EOL>value=error_rate, name='<STR_LIT>', dtype=dtype)<EOL>samples = _check_shape_dominates(samples, [low, high])<EOL>tf.compat.v1.assert_scalar(error_rate)  <EOL>itemwise_error_rate = _itemwise_error_rate(<EOL>total_rate=error_rate, param_tensors=[low, high],<EOL>samples_tensor=samples)<EOL>n = tf.shape(input=samples)[<NUM_LIT:0>]<EOL>envelope = _dkwm_cdf_envelope(n, itemwise_error_rate)<EOL>min_mean = _minimum_mean(samples, envelope, low)<EOL>max_mean = _maximum_mean(samples, envelope, high)<EOL>return min_mean, max_mean<EOL><DEDENT>

Computes a confidence interval for the mean of a scalar distribution.

    In batch mode, computes confidence intervals for all distributions
    in the batch (which need not be identically distributed).

    Relies on the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval).

    The probability (over the randomness of drawing the given samples)
    that any true mean is outside the corresponding returned interval is
    no more than the given `error_rate`.  The size of the intervals
    scale as
    `O(1 / sqrt(#samples))`, as `O(high - low)`, and as `O(-log(error_rate))`.

    Note that `error_rate` is a total error rate for all the confidence
    intervals in the batch.  As such, if the batch is nontrivial, the
    error rate is not broadcast but divided (evenly) among the batch
    members.

    Args:
      samples: Floating-point `Tensor` of samples from the distribution(s)
        of interest.  Entries are assumed IID across the 0th dimension.
        The other dimensions must broadcast with `low` and `high`.
        The support is bounded: `low <= samples <= high`.
      low: Floating-point `Tensor` of lower bounds on the distributions'
        supports.
      high: Floating-point `Tensor` of upper bounds on the distributions'
        supports.
      error_rate: *Scalar* floating-point `Tensor` admissible total rate
        of mistakes.
      name: A name for this operation (optional).

    Returns:
      low: A floating-point `Tensor` of stochastic lower bounds on the
        true means.
      high: A floating-point `Tensor` of stochastic upper bounds on the
        true means.


def _itemwise_error_rate(<EOL>total_rate, param_tensors, samples_tensor=None, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[total_rate, param_tensors, samples_tensor]):<EOL><INDENT>result_shape = [<NUM_LIT:1>]<EOL>for p_tensor in param_tensors:<EOL><INDENT>result_shape = tf.broadcast_dynamic_shape(<EOL>tf.shape(input=p_tensor), result_shape)<EOL><DEDENT>if samples_tensor is not None:<EOL><INDENT>result_shape = tf.broadcast_dynamic_shape(<EOL>tf.shape(input=samples_tensor)[<NUM_LIT:1>:], result_shape)<EOL><DEDENT>num_items = tf.reduce_prod(input_tensor=result_shape)<EOL>return total_rate / tf.cast(num_items, dtype=total_rate.dtype)<EOL><DEDENT>

Distributes a total error rate for a batch of assertions.


def assert_true_mean_equal_by_dkwm(<EOL>samples, low, high, expected, false_fail_rate=<NUM_LIT>, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[samples, low, high, expected, false_fail_rate]):<EOL><INDENT>return assert_true_mean_in_interval_by_dkwm(<EOL>samples, low, high, expected, expected, false_fail_rate)<EOL><DEDENT>

Asserts the mean of the given distribution is as expected.

    More precisely, fails if there is enough evidence (using the
    [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval))
    that the true mean of some distribution from which the given samples are
    drawn is _not_ the given expected mean with statistical significance
    `false_fail_rate` or stronger, otherwise passes.  If you also want to
    check that you are gathering enough evidence that a pass is not
    spurious, see `min_num_samples_for_dkwm_mean_test` and
    `min_discrepancy_of_true_means_detectable_by_dkwm`.

    Note that `false_fail_rate` is a total false failure rate for all
    the assertions in the batch.  As such, if the batch is nontrivial,
    the assertion will insist on stronger evidence to fail any one member.

    Args:
      samples: Floating-point `Tensor` of samples from the distribution(s)
        of interest.  Entries are assumed IID across the 0th dimension.
        The other dimensions must broadcast with `low` and `high`.
        The support is bounded: `low <= samples <= high`.
      low: Floating-point `Tensor` of lower bounds on the distributions'
        supports.
      high: Floating-point `Tensor` of upper bounds on the distributions'
        supports.
      expected: Floating-point `Tensor` of expected true means.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of mistakes.
      name: A name for this operation (optional).

    Returns:
      check: Op that raises `InvalidArgumentError` if any expected mean is
        outside the corresponding confidence interval.


def min_discrepancy_of_true_means_detectable_by_dkwm(<EOL>n, low, high, false_fail_rate, false_pass_rate, name=None):

with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[n, low, high, false_fail_rate, false_pass_rate]):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[n, low, high, false_fail_rate, false_pass_rate], tf.float32)<EOL>n = tf.convert_to_tensor(value=n, name='<STR_LIT:n>', dtype=dtype)<EOL>low = tf.convert_to_tensor(value=low, name='<STR_LIT>', dtype=dtype)<EOL>high = tf.convert_to_tensor(value=high, name='<STR_LIT>', dtype=dtype)<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>false_pass_rate = tf.convert_to_tensor(<EOL>value=false_pass_rate, name='<STR_LIT>', dtype=dtype)<EOL>cdf_discrepancy = min_discrepancy_of_true_cdfs_detectable_by_dkwm(<EOL>n, false_fail_rate, false_pass_rate)<EOL>return (high - low) * cdf_discrepancy<EOL><DEDENT>

Returns the minimum mean discrepancy that a DKWM-based test can detect.

    DKWM is the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval).

    Note that `false_fail_rate` is a total false failure rate for all
    the tests in the batch.  As such, if the batch is nontrivial, each
    member will demand more samples.  The `false_pass_rate` is also
    interpreted as a total, but is treated asymmetrically: If each test
    in the batch detects its corresponding discrepancy with probability
    at least `1 - false_pass_rate`, then running all those tests and
    failing if any one fails will jointly detect all those discrepancies
    with the same `false_pass_rate`.

    Args:
      n: `Tensor` of numbers of samples to be drawn from the distributions
        of interest.
      low: Floating-point `Tensor` of lower bounds on the distributions'
        supports.
      high: Floating-point `Tensor` of upper bounds on the distributions'
        supports.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of false failures.
      false_pass_rate: *Scalar* floating-point `Tensor` admissible rate
        of false passes.
      name: A name for this operation (optional).

    Returns:
      discr: `Tensor` of lower bounds on the distances between true
         means detectable by a DKWM-based test.

    For each batch member `i`, of `K` total, drawing `n[i]` samples from
    some scalar distribution supported on `[low[i], high[i]]` is enough
    to detect a difference in means of size `discr[i]` or more.
    Specifically, we guarantee that (a) if the true mean is the expected
    mean (resp. in the expected interval), then `assert_true_mean_equal_by_dkwm`
    (resp. `assert_true_mean_in_interval_by_dkwm`) will fail with
    probability at most `false_fail_rate / K` (which amounts to
    `false_fail_rate` if applied to the whole batch at once), and (b) if
    the true mean differs from the expected mean (resp. falls outside
    the expected interval) by at least `discr[i]`,
    `assert_true_mean_equal_by_dkwm`
    (resp. `assert_true_mean_in_interval_by_dkwm`) will pass with
    probability at most `false_pass_rate`.

    The detectable discrepancy scales as

    - `O(high[i] - low[i])`,
    - `O(1 / sqrt(n[i]))`,
    - `O(-log(false_fail_rate/K))`, and
    - `O(-log(false_pass_rate))`.


def assert_true_mean_in_interval_by_dkwm(<EOL>samples, low, high, expected_low, expected_high,<EOL>false_fail_rate=<NUM_LIT>, name=None):

args_list = [samples, low, high, expected_low, expected_high, false_fail_rate]<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>args_list):<EOL><INDENT>dtype = dtype_util.common_dtype(args_list, tf.float32)<EOL>samples = tf.convert_to_tensor(value=samples, name='<STR_LIT>', dtype=dtype)<EOL>low = tf.convert_to_tensor(value=low, name='<STR_LIT>', dtype=dtype)<EOL>high = tf.convert_to_tensor(value=high, name='<STR_LIT>', dtype=dtype)<EOL>expected_low = tf.convert_to_tensor(<EOL>value=expected_low, name='<STR_LIT>', dtype=dtype)<EOL>expected_high = tf.convert_to_tensor(<EOL>value=expected_high, name='<STR_LIT>', dtype=dtype)<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>samples = _check_shape_dominates(<EOL>samples, [low, high, expected_low, expected_high])<EOL>min_mean, max_mean = true_mean_confidence_interval_by_dkwm(<EOL>samples, low, high, false_fail_rate)<EOL>check_confidence_interval_can_intersect = tf.compat.v1.assert_greater_equal(<EOL>max_mean,<EOL>expected_low,<EOL>message='<STR_LIT>'<EOL>'<STR_LIT>')<EOL>with tf.control_dependencies([check_confidence_interval_can_intersect]):<EOL><INDENT>return tf.compat.v1.assert_less_equal(<EOL>min_mean,<EOL>expected_high,<EOL>message='<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT><DEDENT>

Asserts the mean of the given distribution is in the given interval.

    More precisely, fails if there is enough evidence (using the
    [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval))
    that the mean of the distribution from which the given samples are
    drawn is _outside_ the given interval with statistical significance
    `false_fail_rate` or stronger, otherwise passes.  If you also want
    to check that you are gathering enough evidence that a pass is not
    spurious, see `min_num_samples_for_dkwm_mean_test` and
    `min_discrepancy_of_true_means_detectable_by_dkwm`.

    Note that `false_fail_rate` is a total false failure rate for all
    the assertions in the batch.  As such, if the batch is nontrivial,
    the assertion will insist on stronger evidence to fail any one member.

    Args:
      samples: Floating-point `Tensor` of samples from the distribution(s)
        of interest.  Entries are assumed IID across the 0th dimension.
        The other dimensions must broadcast with `low` and `high`.
        The support is bounded: `low <= samples <= high`.
      low: Floating-point `Tensor` of lower bounds on the distributions'
        supports.
      high: Floating-point `Tensor` of upper bounds on the distributions'
        supports.
      expected_low: Floating-point `Tensor` of lower bounds on the
        expected true means.
      expected_high: Floating-point `Tensor` of upper bounds on the
        expected true means.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of mistakes.
      name: A name for this operation (optional).

    Returns:
      check: Op that raises `InvalidArgumentError` if any expected mean
        interval does not overlap with the corresponding confidence
        interval.


def assert_true_mean_equal_by_dkwm_two_sample(<EOL>samples1, low1, high1, samples2, low2, high2,<EOL>false_fail_rate=<NUM_LIT>, name=None):

args_list = [samples1, low1, high1, samples2, low2, high2, false_fail_rate]<EOL>with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', args_list):<EOL><INDENT>dtype = dtype_util.common_dtype(args_list, tf.float32)<EOL>samples1 = tf.convert_to_tensor(<EOL>value=samples1, name='<STR_LIT>', dtype=dtype)<EOL>low1 = tf.convert_to_tensor(value=low1, name='<STR_LIT>', dtype=dtype)<EOL>high1 = tf.convert_to_tensor(value=high1, name='<STR_LIT>', dtype=dtype)<EOL>samples2 = tf.convert_to_tensor(<EOL>value=samples2, name='<STR_LIT>', dtype=dtype)<EOL>low2 = tf.convert_to_tensor(value=low2, name='<STR_LIT>', dtype=dtype)<EOL>high2 = tf.convert_to_tensor(value=high2, name='<STR_LIT>', dtype=dtype)<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>samples1 = _check_shape_dominates(samples1, [low1, high1])<EOL>samples2 = _check_shape_dominates(samples2, [low2, high2])<EOL>compatible_samples = tf.compat.v1.assert_equal(<EOL>tf.shape(input=samples1)[<NUM_LIT:1>:],<EOL>tf.shape(input=samples2)[<NUM_LIT:1>:])<EOL>with tf.control_dependencies([compatible_samples]):<EOL><INDENT>min_mean_2, max_mean_2 = true_mean_confidence_interval_by_dkwm(<EOL>samples2, low2, high2, false_fail_rate / <NUM_LIT>)<EOL>return assert_true_mean_in_interval_by_dkwm(<EOL>samples1, low1, high1, min_mean_2, max_mean_2, false_fail_rate / <NUM_LIT>)<EOL><DEDENT><DEDENT>

Asserts the means of the given distributions are equal.

    More precisely, fails if there is enough evidence (using the
    [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval))
    that the means of the distributions from which the given samples are
    drawn are _not_ equal with statistical significance `false_fail_rate`
    or stronger, otherwise passes.  If you also want to check that you
    are gathering enough evidence that a pass is not spurious, see
    `min_num_samples_for_dkwm_mean_two_sample_test` and
    `min_discrepancy_of_true_means_detectable_by_dkwm_two_sample`.

    Note that `false_fail_rate` is a total false failure rate for all
    the assertions in the batch.  As such, if the batch is nontrivial,
    the assertion will insist on stronger evidence to fail any one member.

    Args:
      samples1: Floating-point `Tensor` of samples from the
        distribution(s) A.  Entries are assumed IID across the 0th
        dimension.  The other dimensions must broadcast with `low1`,
        `high1`, `low2`, and `high2`.
        The support is bounded: `low1 <= samples1 <= high1`.
      low1: Floating-point `Tensor` of lower bounds on the supports of the
        distributions A.
      high1: Floating-point `Tensor` of upper bounds on the supports of
        the distributions A.
      samples2: Floating-point `Tensor` of samples from the
        distribution(s) B.  Entries are assumed IID across the 0th
        dimension.  The other dimensions must broadcast with `low1`,
        `high1`, `low2`, and `high2`.
        The support is bounded: `low2 <= samples2 <= high2`.
      low2: Floating-point `Tensor` of lower bounds on the supports of the
        distributions B.
      high2: Floating-point `Tensor` of upper bounds on the supports of
        the distributions B.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of mistakes.
      name: A name for this operation (optional).

    Returns:
      check: Op that raises `InvalidArgumentError` if any pair of confidence
        intervals true for corresponding true means do not overlap.


def min_discrepancy_of_true_means_detectable_by_dkwm_two_sample(<EOL>n1, low1, high1, n2, low2, high2,<EOL>false_fail_rate, false_pass_rate, name=None):

args_list = (<EOL>[n1, low1, high1, n2, low2, high2, false_fail_rate, false_pass_rate])<EOL>with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>args_list):<EOL><INDENT>dtype = dtype_util.common_dtype(args_list, tf.float32)<EOL>n1 = tf.convert_to_tensor(value=n1, name='<STR_LIT>', dtype=dtype)<EOL>low1 = tf.convert_to_tensor(value=low1, name='<STR_LIT>', dtype=dtype)<EOL>high1 = tf.convert_to_tensor(value=high1, name='<STR_LIT>', dtype=dtype)<EOL>n2 = tf.convert_to_tensor(value=n2, name='<STR_LIT>', dtype=dtype)<EOL>low2 = tf.convert_to_tensor(value=low2, name='<STR_LIT>', dtype=dtype)<EOL>high2 = tf.convert_to_tensor(value=high2, name='<STR_LIT>', dtype=dtype)<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>false_pass_rate = tf.convert_to_tensor(<EOL>value=false_pass_rate, name='<STR_LIT>', dtype=dtype)<EOL>det_disc1 = min_discrepancy_of_true_means_detectable_by_dkwm(<EOL>n1, low1, high1, false_fail_rate / <NUM_LIT>, false_pass_rate / <NUM_LIT>)<EOL>det_disc2 = min_discrepancy_of_true_means_detectable_by_dkwm(<EOL>n2, low2, high2, false_fail_rate / <NUM_LIT>, false_pass_rate / <NUM_LIT>)<EOL>return det_disc1 + det_disc2<EOL><DEDENT>

Returns the minimum mean discrepancy for a two-sample DKWM-based test.

    DKWM is the [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
    (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval).

    Note that `false_fail_rate` is a total false failure rate for all
    the tests in the batch.  As such, if the batch is nontrivial, each
    member will demand more samples.  The `false_pass_rate` is also
    interpreted as a total, but is treated asymmetrically: If each test
    in the batch detects its corresponding discrepancy with probability
    at least `1 - false_pass_rate`, then running all those tests and
    failing if any one fails will jointly detect all those discrepancies
    with the same `false_pass_rate`.

    Args:
      n1: `Tensor` of numbers of samples to be drawn from the distributions A.
      low1: Floating-point `Tensor` of lower bounds on the supports of the
        distributions A.
      high1: Floating-point `Tensor` of upper bounds on the supports of
        the distributions A.
      n2: `Tensor` of numbers of samples to be drawn from the distributions B.
      low2: Floating-point `Tensor` of lower bounds on the supports of the
        distributions B.
      high2: Floating-point `Tensor` of upper bounds on the supports of
        the distributions B.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of false failures.
      false_pass_rate: *Scalar* floating-point `Tensor` admissible rate
        of false passes.
      name: A name for this operation (optional).

    Returns:
      discr: `Tensor` of lower bounds on the distances between true means
         detectable by a two-sample DKWM-based test.

    For each batch member `i`, of `K` total, drawing `n1[i]` samples
    from scalar distribution A supported on `[low1[i], high1[i]]` and `n2[i]`
    samples from scalar distribution B supported on `[low2[i], high2[i]]`
    is enough to detect a difference in their true means of size
    `discr[i]` or more.  Specifically, we guarantee that (a) if their
    true means are equal, `assert_true_mean_equal_by_dkwm_two_sample`
    will fail with probability at most `false_fail_rate/K` (which
    amounts to `false_fail_rate` if applied to the whole batch at once),
    and (b) if their true means differ by at least `discr[i]`,
    `assert_true_mean_equal_by_dkwm_two_sample` will pass with
    probability at most `false_pass_rate`.

    The detectable distribution scales as

    - `O(high1[i] - low1[i])`, `O(high2[i] - low2[i])`,
    - `O(1 / sqrt(n1[i]))`, `O(1 / sqrt(n2[i]))`,
    - `O(-log(false_fail_rate/K))`, and
    - `O(-log(false_pass_rate))`.


def assert_multivariate_true_cdf_equal_on_projections_two_sample(<EOL>samples1, samples2, num_projections, event_ndims=<NUM_LIT:1>,<EOL>false_fail_rate=<NUM_LIT>, seed=<NUM_LIT>, name=None):

<EOL>args_list = (<EOL>[samples1, samples2, num_projections, event_ndims, false_fail_rate])<EOL>strm = seed_stream.SeedStream(salt='<STR_LIT>', seed=seed)<EOL>with tf.compat.v1.name_scope(<EOL>name,<EOL>'<STR_LIT>',<EOL>args_list):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[samples1, samples2, false_fail_rate], tf.float32)<EOL>samples1 = tf.convert_to_tensor(<EOL>value=samples1, name='<STR_LIT>', dtype=dtype)<EOL>samples2 = tf.convert_to_tensor(<EOL>value=samples2, name='<STR_LIT>', dtype=dtype)<EOL>num_projections = tf.convert_to_tensor(<EOL>value=num_projections, name='<STR_LIT>')<EOL>false_fail_rate = tf.convert_to_tensor(<EOL>value=false_fail_rate, name='<STR_LIT>', dtype=dtype)<EOL>tf.compat.v1.assert_scalar(false_fail_rate)  <EOL>compatible_samples = tf.compat.v1.assert_equal(<EOL>tf.shape(input=samples1)[<NUM_LIT:1>:],<EOL>tf.shape(input=samples2)[<NUM_LIT:1>:])<EOL>with tf.control_dependencies([compatible_samples]):<EOL><INDENT>event_shape = tf.shape(input=samples1)[-event_ndims:]<EOL>random_projections = _random_unit_hypersphere(<EOL>[num_projections], event_shape, dtype=dtype, seed=strm())<EOL>last_axes = list(range(-<NUM_LIT:1>, -(event_ndims+<NUM_LIT:1>), -<NUM_LIT:1>))<EOL>proj1 = tf.tensordot(samples1, random_projections, [last_axes, last_axes])<EOL>proj2 = tf.tensordot(samples2, random_projections, [last_axes, last_axes])<EOL>return assert_true_cdf_equal_by_dkwm_two_sample(<EOL>proj1, proj2, false_fail_rate=false_fail_rate)<EOL><DEDENT><DEDENT>

Asserts the given multivariate distributions are equal.

    The test is a 1-D DKWM-style test of equality in distribution along the given
    number of random projections.  This is of course imperfect, but can behave
    reasonably, especially if `num_projections` is significantly more than the
    dimensionality of the sample space.

    More precisely, the test
    (i) assumes the event shape is given by the trailing `event_ndims` dimensions
        in each `samples` Tensor;
    (ii) generates `num_projections` random projections from this space to scalar;
    (iii) fails if there is enough evidence (using the
        [Dvoretzky-Kiefer-Wolfowitz-Massart inequality]
        (https://en.wikipedia.org/wiki/CDF-based_nonparametric_confidence_interval))
        along any of these projections that `samples1` and `samples2` come from
        different true distributions.

    This test works as written even if the distribution in question has atoms
    (e.g., is discrete).

    Note that the top dimension of each `samples` is treated as iid, and the
    bottom `event_ndims` dimensions are projected to scalar.  The remaining
    dimensions, if any, are treated as batch dimensions, and `false_fail_rate` is
    a total false failure rate for all the assertions in the batch (and all
    projections).  As such, if the batch is nontrivial, the assertion will insist
    on stronger evidence to fail any one member.

    A note on experiment design: This test boils down to `num_projections`
    two-sample CDF equality tests.  As such, one can compute the number of samples
    to draw or the detectable discrepancy (along any projection) using
    `min_num_samples_for_dkwm_cdf_two_sample_test` and
    `min_discrepancy_of_true_cdfs_detectable_by_dkwm_two_sample` respectively,
    just being sure to divide the false failure rate by the number of projections
    requested (no need to adjust the false pass rate).

    Args:
      samples1: Tensor of shape [n] + B + E.  Samples from some (batch of)
        distribution(s) of interest.  Assumed IID across the 0 dimension.
      samples2: Tensor of shape [m] + B + E.  Samples from some (batch of)
        distribution(s) of interest.  Assumed IID across the 0 dimension.
      num_projections: Scalar integer Tensor.  Number of projections to use.  Each
        projection will be a random direction on the unit hypershpere of shape E.
      event_ndims: Number of trailing dimensions forming the event shape.
        `rank(E) == event_ndims`.
      false_fail_rate: *Scalar* floating-point `Tensor` admissible total
        rate of mistakes.
      seed: Optional PRNG seed to use for generating the random projections.
        Changing this from the default should generally not be necessary.
      name: A name for this operation (optional).

    Returns:
      check: Op that raises `InvalidArgumentError` if the two samples do not match
        along any of the generated projections.


def assign_moving_mean_variance(<EOL>mean_var, variance_var, value, decay, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>",<EOL>[variance_var, mean_var, value, decay]):<EOL><INDENT>with tf.compat.v1.colocate_with(variance_var):<EOL><INDENT>with tf.compat.v1.colocate_with(mean_var):<EOL><INDENT>base_dtype = mean_var.dtype.base_dtype<EOL>if not base_dtype.is_floating:<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(base_dtype.name))<EOL><DEDENT>if base_dtype != variance_var.dtype.base_dtype:<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>".format(<EOL>base_dtype.name,<EOL>variance_var.dtype.base_dtype.name))<EOL><DEDENT>value = tf.convert_to_tensor(<EOL>value=value, dtype=base_dtype, name="<STR_LIT:value>")<EOL>decay = tf.convert_to_tensor(<EOL>value=decay, dtype=base_dtype, name="<STR_LIT>")<EOL>delta = value - mean_var<EOL>with tf.control_dependencies([delta]):<EOL><INDENT>mean_var = mean_var.assign_add((<NUM_LIT:1.> - decay) * delta)<EOL>variance_var = variance_var.assign_sub(<EOL>(<NUM_LIT:1.> - decay) * (variance_var - decay * tf.square(delta)))<EOL><DEDENT>return mean_var, variance_var<EOL><DEDENT><DEDENT><DEDENT>

Compute exponentially weighted moving {mean,variance} of a streaming value.

    The `value` updated exponentially weighted moving `mean_var` and
    `variance_var` are given by the following recurrence relations:

    ```python
    variance_var = decay * (variance_var + (1 - decay) * (value - mean_var)**2)
    mean_var     = decay * mean_var + (1 - decay) * value
    ```

    Note: `mean_var` is updated *after* `variance_var`, i.e., `variance_var` uses
    the lag-1 mean.

    For derivation justification, see [Finch (2009; Eq. 143)][1].
    Parameterization: Finch's `alpha` is `1 - decay`.

    Args:
      mean_var: `float`-like `Variable` representing the exponentially weighted
        moving mean. Same shape as `variance_var` and `value`.
      variance_var: `float`-like `Variable` representing the
        exponentially weighted moving variance. Same shape as `mean_var` and
        `value`.
      value: `float`-like `Tensor`. Same shape as `mean_var` and `variance_var`.
      decay: A `float`-like `Tensor`. The moving mean decay. Typically close to
        `1.`, e.g., `0.999`.
      name: Optional name of the returned operation.

    Returns:
      mean_var: `Variable` representing the `value`-updated exponentially weighted
        moving mean.
      variance_var: `Variable` representing the `value`-updated
        exponentially weighted moving variance.

    Raises:
      TypeError: if `mean_var` does not have float type `dtype`.
      TypeError: if `mean_var`, `variance_var`, `value`, `decay` have different
        `base_dtype`.

    #### References

    [1]: Tony Finch. Incremental calculation of weighted mean and variance.
         _Technical Report_, 2009.
         http://people.ds.cam.ac.uk/fanf2/hermes/doc/antiforgery/stats.pdf


def assign_log_moving_mean_exp(<EOL>log_mean_exp_var, log_value, decay, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>",<EOL>[log_mean_exp_var, log_value, decay]):<EOL><INDENT>with tf.compat.v1.colocate_with(log_mean_exp_var):<EOL><INDENT>base_dtype = log_mean_exp_var.dtype.base_dtype<EOL>if not base_dtype.is_floating:<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(base_dtype.name))<EOL><DEDENT>log_value = tf.convert_to_tensor(<EOL>value=log_value, dtype=base_dtype, name="<STR_LIT>")<EOL>decay = tf.convert_to_tensor(value=decay, dtype=base_dtype, name="<STR_LIT>")<EOL>delta = (log_value - log_mean_exp_var)[tf.newaxis, ...]<EOL>x = tf.concat([<EOL>tf.math.log(decay) * tf.ones_like(delta),<EOL>delta + tf.math.log1p(-decay)<EOL>],<EOL>axis=<NUM_LIT:0>)<EOL>x = tf.reduce_logsumexp(input_tensor=x, axis=<NUM_LIT:0>)<EOL>return log_mean_exp_var.assign_add(x)<EOL><DEDENT><DEDENT>

Compute the log of the exponentially weighted moving mean of the exp.

    If `log_value` is a draw from a stationary random variable, this function
    approximates `log(E[exp(log_value)])`, i.e., a weighted log-sum-exp. More
    precisely, a `tf.Variable`, `log_mean_exp_var`, is updated by `log_value`
    using the following identity:

    ```none
    log_mean_exp_var =
    = log(decay exp(log_mean_exp_var) + (1 - decay) exp(log_value))
    = log(exp(log_mean_exp_var + log(decay)) + exp(log_value + log1p(-decay)))
    = log_mean_exp_var
      + log(  exp(log_mean_exp_var   - log_mean_exp_var + log(decay))
            + exp(log_value - log_mean_exp_var + log1p(-decay)))
    = log_mean_exp_var
      + log_sum_exp([log(decay), log_value - log_mean_exp_var + log1p(-decay)]).
    ```

    In addition to numerical stability, this formulation is advantageous because
    `log_mean_exp_var` can be updated in a lock-free manner, i.e., using
    `assign_add`. (Note: the updates are not thread-safe; it's just that the
    update to the tf.Variable is presumed efficient due to being lock-free.)

    Args:
      log_mean_exp_var: `float`-like `Variable` representing the log of the
        exponentially weighted moving mean of the exp. Same shape as `log_value`.
      log_value: `float`-like `Tensor` representing a new (streaming) observation.
        Same shape as `log_mean_exp_var`.
      decay: A `float`-like `Tensor`. The moving mean decay. Typically close to
        `1.`, e.g., `0.999`.
      name: Optional name of the returned operation.

    Returns:
      log_mean_exp_var: A reference to the input 'Variable' tensor with the
        `log_value`-updated log of the exponentially weighted moving mean of exp.

    Raises:
      TypeError: if `log_mean_exp_var` does not have float type `dtype`.
      TypeError: if `log_mean_exp_var`, `log_value`, `decay` have different
        `base_dtype`.


@deprecation.deprecated(<EOL>"<STR_LIT>", "<STR_LIT>"<EOL>"<STR_LIT>")<EOL>def moving_mean_variance(value, decay, name=None):

with tf.compat.v1.variable_scope(name, "<STR_LIT>",<EOL>[value, decay]):<EOL><INDENT>value = tf.convert_to_tensor(value=value, name="<STR_LIT:value>")<EOL>base_dtype = value.dtype.base_dtype<EOL>if not base_dtype.is_floating:<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>".format(<EOL>base_dtype.name))<EOL><DEDENT>decay = tf.convert_to_tensor(value=decay, dtype=base_dtype, name="<STR_LIT>")<EOL>variance_var = tf.compat.v2.Variable(<EOL>name="<STR_LIT>",<EOL>initial_value=tf.zeros(shape=value.shape, dtype=value.dtype),<EOL>trainable=False)<EOL>mean_var = tf.compat.v2.Variable(<EOL>name="<STR_LIT>",<EOL>initial_value=tf.zeros(shape=value.shape, dtype=value.dtype),<EOL>trainable=False)<EOL>return assign_moving_mean_variance(<EOL>mean_var, variance_var, value, decay)<EOL><DEDENT>

Compute exponentially weighted moving {mean,variance} of a streaming value.

    The exponentially-weighting moving `mean_var` and `variance_var` are updated
    by `value` according to the following recurrence:

    ```python
    variance_var = decay * (variance_var + (1-decay) * (value - mean_var)**2)
    mean_var     = decay * mean_var + (1 - decay) * value
    ```

    Note: `mean_var` is updated *after* `variance_var`, i.e., `variance_var` uses
    the lag-`1` mean.

    For derivation justification, see [Finch (2009; Eq. 143)][1].

    Unlike `assign_moving_mean_variance`, this function handles
    variable creation.

    Args:
      value: `float`-like `Tensor`. Same shape as `mean_var` and `variance_var`.
      decay: A `float`-like `Tensor`. The moving mean decay. Typically close to
        `1.`, e.g., `0.999`.
      name: Optional name of the returned operation.

    Returns:
      mean_var: `Variable` representing the `value`-updated exponentially weighted
        moving mean.
      variance_var: `Variable` representing the `value`-updated
        exponentially weighted moving variance.

    Raises:
      TypeError: if `value_var` does not have float type `dtype`.
      TypeError: if `value`, `decay` have different `base_dtype`.

    #### References

    [1]: Tony Finch. Incremental calculation of weighted mean and variance.
         _Technical Report_, 2009.
         http://people.ds.cam.ac.uk/fanf2/hermes/doc/antiforgery/stats.pdf


@kullback_leibler.RegisterKL(HalfNormal, HalfNormal)<EOL>def _kl_half_normal_half_normal(a, b, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>return (tf.math.log(b.scale) - tf.math.log(a.scale) +<EOL>(a.scale**<NUM_LIT:2> - b.scale**<NUM_LIT:2>) / (<NUM_LIT:2> * b.scale**<NUM_LIT:2>))<EOL><DEDENT>

Calculate the batched KL divergence KL(a || b) with a and b `HalfNormal`.

    Args:
      a: Instance of a `HalfNormal` distribution object.
      b: Instance of a `HalfNormal` distribution object.
      name: (optional) Name to use for created operations.
        default is "kl_half_normal_half_normal".

    Returns:
      Batchwise KL(a || b)


parameters = dict(locals())<EOL>with tf.name_scope(name) as name:<EOL><INDENT>scale = tf.convert_to_tensor(<EOL>value=scale,<EOL>name="<STR_LIT>",<EOL>dtype=dtype_util.common_dtype([scale], preferred_dtype=tf.float32))<EOL>with tf.control_dependencies(<EOL>[assert_util.assert_positive(scale)] if validate_args else []):<EOL><INDENT>self._scale = tf.identity(scale, name="<STR_LIT>")<EOL><DEDENT><DEDENT>super(HalfNormal, self).__init__(<EOL>dtype=self._scale.dtype,<EOL>reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>parameters=parameters,<EOL>graph_parents=[self._scale],<EOL>name=name)<EOL>

Construct HalfNormals with scale `scale`.

        Args:
          scale: Floating point tensor; the scales of the distribution(s).
            Must contain only positive values.
          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: Python `str` name prefixed to Ops created by this class.


return self._scale<EOL>


def _compute_non_batch_kl(self, loc_a, scale_a, loc_b, scale_b):

<EOL>covariance_a = np.dot(scale_a, scale_a.T)<EOL>covariance_b = np.dot(scale_b, scale_b.T)<EOL>covariance_b_inv = np.linalg.inv(covariance_b)<EOL>t = np.trace(covariance_b_inv.dot(covariance_a))<EOL>q = (loc_b - loc_a).dot(covariance_b_inv).dot(loc_b - loc_a)<EOL>k = loc_a.shape[<NUM_LIT:0>]<EOL>l = np.log(np.linalg.det(covariance_b) / np.linalg.det(covariance_a))<EOL>return <NUM_LIT:0.5> * (t + q - k + l)<EOL>

Non-batch KL for N(loc_a, scale_a), N(loc_b, scale_b).


@hps.composite<EOL>def broadcasting_shapes(draw, batch_shape, param_names):

n = len(param_names)<EOL>return dict(zip(draw(hps.permutations(param_names)),<EOL>draw(tfp_test_util.broadcasting_shapes(batch_shape, n))))<EOL>

Draws a set of parameter batch shapes that broadcast to `batch_shape`.

    For each parameter we need to choose its batch rank, and whether or not each
    axis i is 1 or batch_shape[i]. This function chooses a set of shapes that
    have possibly mismatched ranks, and possibly broadcasting axes, with the
    promise that the broadcast of the set of all shapes matches `batch_shape`.

    Args:
      draw: Hypothesis sampler.
      batch_shape: `tf.TensorShape`, the target (fully-defined) batch shape .
      param_names: Iterable of `str`, the parameters whose batch shapes need
        determination.

    Returns:
      param_batch_shapes: `dict` of `str->tf.TensorShape` where the set of
          shapes broadcast to `batch_shape`. The shapes are fully defined.