def matrix_rank(a, tol=None, validate_args=False, name=None): """Compute the matrix rank; the number of non-zero SVD singular values. Arguments: a: (Batch of) `float`-like matrix-shaped `Tensor`(s) which are to be pseudo-inverted. tol: Threshold below which the singular value is counted as 'zero'. Default value: `None` (i.e., `eps * max(rows, cols) * max(singular_val)`). validate_args: When `True`, additional assertions might be embedded in the graph. Default value: `False` (i.e., no graph assertions are added). name: Python `str` prefixed to ops created by this function. Default value: 'matrix_rank'. Returns: matrix_rank: (Batch of) `int32` scalars representing the number of non-zero singular values. """ with tf.name_scope(name or 'matrix_rank'): a = tf.convert_to_tensor(a, dtype_hint=tf.float32, name='a') assertions = _maybe_validate_matrix(a, validate_args) if assertions: with tf.control_dependencies(assertions): a = tf.identity(a) s = tf.linalg.svd(a, compute_uv=False) if tol is None: if tensorshape_util.is_fully_defined(a.shape[-2:]): m = np.max(a.shape[-2:].as_list()) else: m = tf.reduce_max(tf.shape(a)[-2:]) eps = np.finfo(dtype_util.as_numpy_dtype(a.dtype)).eps tol = (eps * tf.cast(m, a.dtype) * tf.reduce_max(s, axis=-1, keepdims=True)) return tf.reduce_sum(tf.cast(s > tol, tf.int32), axis=-1)
def _make_perm(self, x_rank, perm): sample_batch_ndims = (distribution_util.prefer_static_value(x_rank) - distribution_util.prefer_static_value( self.rightmost_transposed_ndims)) dtype = perm.dtype perm = tf.concat([ tf.range(tf.cast(sample_batch_ndims, dtype)), tf.cast( sample_batch_ndims + distribution_util.prefer_static_value(perm), dtype), ], axis=0) return perm
def _prob(self, event): samples = tf.convert_to_tensor(self._samples) num_samples = self._compute_num_samples(samples) event = tf.convert_to_tensor(event, name='event', dtype=self.dtype) event, samples = _broadcast_event_and_samples(event, samples, event_ndims=self._event_ndims) prob = tf.reduce_sum( tf.cast( tf.reduce_all( tf.equal(samples, event), axis=tf.range(-self._event_ndims, 0)), dtype=tf.int32), axis=-1) / num_samples if dtype_util.is_floating(self.dtype): prob = tf.cast(prob, self.dtype) return prob
def _cdf(self, k): # TODO(b/135263541): Improve numerical precision of categorical.cdf. probs = self.probs_parameter() num_categories = self._num_categories(probs) k, probs = _broadcast_cat_event_and_params( k, probs, base_dtype=dtype_util.base_dtype(self.dtype)) # Since the lowest number in the support is 0, any k < 0 should be zero in # the output. should_be_zero = k < 0 # Will use k as an index in the gather below, so clip it to {0,...,K-1}. k = tf.clip_by_value(tf.cast(k, tf.int32), 0, num_categories - 1) batch_shape = tf.shape(k) # tf.gather(..., batch_dims=batch_dims) requires static batch_dims kwarg, so # to handle the case where the batch shape is dynamic, flatten the batch # dims (so we know batch_dims=1). k_flat_batch = tf.reshape(k, [-1]) probs_flat_batch = tf.reshape( probs, tf.concat(([-1], [num_categories]), axis=0)) cdf_flat = tf.gather(tf.cumsum(probs_flat_batch, axis=-1), k_flat_batch[..., tf.newaxis], batch_dims=1) cdf = tf.reshape(cdf_flat, shape=batch_shape) zero = np.array(0, dtype=dtype_util.as_numpy_dtype(cdf.dtype)) return tf.where(should_be_zero, zero, cdf)
def _multi_gamma_sequence(self, a, p, name="multi_gamma_sequence"): """Creates sequence used in multivariate (di)gamma; shape = shape(a)+[p].""" with self._name_and_control_scope(name): # Linspace only takes scalars, so we'll add in the offset afterwards. seq = tf.linspace(tf.constant(0., dtype=self.dtype), 0.5 - 0.5 * p, tf.cast(p, tf.int32)) return seq + tf.expand_dims(a, [-1])
def reduce_logmeanexp(input_tensor, axis=None, keepdims=False, name=None): """Computes `log(mean(exp(input_tensor)))`. Reduces `input_tensor` along the dimensions given in `axis`. Unless `keepdims` is true, the rank of the tensor is reduced by 1 for each entry in `axis`. If `keepdims` is true, the reduced dimensions are retained with length 1. If `axis` has no entries, all dimensions are reduced, and a tensor with a single element is returned. This function is more numerically stable than `log(reduce_mean(exp(input)))`. It avoids overflows caused by taking the exp of large inputs and underflows caused by taking the log of small inputs. Args: input_tensor: The tensor to reduce. Should have numeric type. axis: The dimensions to reduce. If `None` (the default), reduces all dimensions. Must be in the range `[-rank(input_tensor), rank(input_tensor))`. keepdims: Boolean. Whether to keep the axis as singleton dimensions. Default value: `False` (i.e., squeeze the reduced dimensions). name: Python `str` name prefixed to Ops created by this function. Default value: `None` (i.e., `'reduce_logmeanexp'`). Returns: log_mean_exp: The reduced tensor. """ with tf.name_scope(name or 'reduce_logmeanexp'): lse = tf.reduce_logsumexp(input_tensor, axis=axis, keepdims=keepdims) n = prefer_static.size(input_tensor) // prefer_static.size(lse) log_n = tf.math.log(tf.cast(n, lse.dtype)) return lse - log_n
def _mode(self, samples=None): # Samples count can vary by batch member. Use map_fn to compute mode for # each batch separately. def _get_mode(samples): # TODO(b/123985779): Switch to tf.unique_with_counts_v2 when exposed count = gen_array_ops.unique_with_counts_v2(samples, axis=[0]).count return tf.argmax(count) if samples is None: samples = tf.convert_to_tensor(self._samples) num_samples = self._compute_num_samples(samples) # Flatten samples for each batch. if self._event_ndims == 0: flattened_samples = tf.reshape(samples, [-1, num_samples]) mode_shape = self._batch_shape_tensor(samples) else: event_size = tf.reduce_prod(self._event_shape_tensor(samples)) mode_shape = tf.concat( [self._batch_shape_tensor(samples), self._event_shape_tensor(samples)], axis=0) flattened_samples = tf.reshape(samples, [-1, num_samples, event_size]) indices = tf.map_fn(_get_mode, flattened_samples, dtype=tf.int64) full_indices = tf.stack( [tf.range(tf.shape(indices)[0]), tf.cast(indices, tf.int32)], axis=1) mode = tf.gather_nd(flattened_samples, full_indices) return tf.reshape(mode, mode_shape)
def _log_prob(self, x): logits = self._logits_parameter_no_checks() event_size = self._event_size(logits) x = tf.cast(x, logits.dtype) x = self._maybe_assert_valid_sample(x, dtype=logits.dtype) # broadcast logits or x if need be. if (not tensorshape_util.is_fully_defined(x.shape) or not tensorshape_util.is_fully_defined(logits.shape) or x.shape != logits.shape): broadcast_shape = tf.broadcast_dynamic_shape( tf.shape(logits), tf.shape(x)) logits = tf.broadcast_to(logits, broadcast_shape) x = tf.broadcast_to(x, broadcast_shape) logits_shape = tf.shape(tf.reduce_sum(logits, axis=-1)) logits_2d = tf.reshape(logits, [-1, event_size]) x_2d = tf.reshape(x, [-1, event_size]) ret = -tf.nn.softmax_cross_entropy_with_logits( labels=tf.stop_gradient(x_2d), logits=logits_2d) # Reshape back to user-supplied batch and sample dims prior to 2D reshape. ret = tf.reshape(ret, logits_shape) return ret
def _log_prob(self, x, power=None): # The log probability at positive integer points x is log(x^(-power) / Z) # where Z is the normalization constant. For x < 1 and non-integer points, # the log-probability is -inf. # # However, if interpolate_nondiscrete is True, we return the natural # continuous relaxation for x >= 1 which agrees with the log probability at # positive integer points. # # If interpolate_nondiscrete is False and validate_args is True, we check # that the sample point x is in the support. That is, x is equivalent to a # positive integer. power = power if power is not None else tf.convert_to_tensor( self.power) x = tf.cast(x, power.dtype) if self.validate_args and not self.interpolate_nondiscrete: x = distribution_util.embed_check_integer_casting_closed( x, target_dtype=self.dtype, assert_positive=True) log_normalization = tf.math.log(tf.math.zeta(power, 1.)) safe_x = tf.maximum(x if self.interpolate_nondiscrete else tf.floor(x), 1.) y = -power * tf.math.log(safe_x) log_unnormalized_prob = tf.where( tf.equal(x, safe_x), y, dtype_util.as_numpy_dtype(y.dtype)(-np.inf)) return log_unnormalized_prob - log_normalization
def _sample_n(self, n, seed=None): n_draws = tf.cast(self.total_count, dtype=tf.int32) logits = self._logits_parameter_no_checks() k = tf.compat.dimension_value(logits.shape[-1]) if k is None: k = tf.shape(logits)[-1] return draw_sample(n, k, logits, n_draws, self.dtype, seed)
def _entropy(self): concentration = tf.convert_to_tensor(self.concentration) k = tf.cast(tf.shape(concentration)[-1], self.dtype) total_concentration = tf.reduce_sum(concentration, axis=-1) return (tf.math.lbeta(concentration) + ((total_concentration - k) * tf.math.digamma(total_concentration)) - tf.reduce_sum((concentration - 1.) * tf.math.digamma(concentration), axis=-1))
def _log_prob(self, event): if self.validate_args: event = distribution_util.embed_check_integer_casting_closed( event, target_dtype=tf.bool) log_probs0, log_probs1 = self._outcome_log_probs() event = tf.cast(event, log_probs0.dtype) return event * (log_probs1 - log_probs0) + log_probs0
def _entropy(self, **kwargs): if not self.bijector.is_constant_jacobian: raise NotImplementedError("entropy is not implemented") if not self.bijector._is_injective: # pylint: disable=protected-access raise NotImplementedError("entropy is not implemented when " "bijector is not injective.") distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs) # Suppose Y = g(X) where g is a diffeomorphism and X is a continuous rv. It # can be shown that: # H[Y] = H[X] + E_X[(log o abs o det o J o g)(X)]. # If is_constant_jacobian then: # E_X[(log o abs o det o J o g)(X)] = (log o abs o det o J o g)(c) # where c can by anything. entropy = self.distribution.entropy(**distribution_kwargs) if self._is_maybe_event_override: # H[X] = sum_i H[X_i] if X_i are mutually independent. # This means that a reduce_sum is a simple rescaling. entropy = entropy * tf.cast( tf.reduce_prod(self._override_event_shape), dtype=dtype_util.base_dtype(entropy.dtype)) if self._is_maybe_batch_override: new_shape = tf.concat([ prefer_static.ones_like(self._override_batch_shape), self.distribution.batch_shape_tensor() ], 0) entropy = tf.reshape(entropy, new_shape) multiples = tf.concat([ self._override_batch_shape, prefer_static.ones_like(self.distribution.batch_shape_tensor()) ], 0) entropy = tf.tile(entropy, multiples) dummy = prefer_static.zeros(shape=tf.concat( [self.batch_shape_tensor(), self.event_shape_tensor()], 0), dtype=self.dtype) event_ndims = (tensorshape_util.rank(self.event_shape) if tensorshape_util.rank(self.event_shape) is not None else tf.size(self.event_shape_tensor())) ildj = self.bijector.inverse_log_det_jacobian(dummy, event_ndims=event_ndims, **bijector_kwargs) entropy = entropy - tf.cast(ildj, entropy.dtype) tensorshape_util.set_shape(entropy, self.batch_shape) return entropy
def body(m, pchol, perm, matrix_diag): """Body of a single `tf.while_loop` iteration.""" # Here is roughly a numpy, non-batched version of what's going to happen. # (See also Algorithm 1 of Harbrecht et al.) # 1: maxi = np.argmax(matrix_diag[perm[m:]]) + m # 2: maxval = matrix_diag[perm][maxi] # 3: perm[m], perm[maxi] = perm[maxi], perm[m] # 4: row = matrix[perm[m]][perm[m + 1:]] # 5: row -= np.sum(pchol[:m][perm[m + 1:]] * pchol[:m][perm[m]]], axis=-2) # 6: pivot = np.sqrt(maxval); row /= pivot # 7: row = np.concatenate([[[pivot]], row], -1) # 8: matrix_diag[perm[m:]] -= row**2 # 9: pchol[m, perm[m:]] = row # Find the maximal position of the (remaining) permuted diagonal. # Steps 1, 2 above. permuted_diag = batch_gather(matrix_diag, perm[..., m:]) maxi = tf.argmax(permuted_diag, axis=-1, output_type=tf.int64)[..., tf.newaxis] maxval = batch_gather(permuted_diag, maxi) maxi = maxi + m maxval = maxval[..., 0] # Update perm: Swap perm[...,m] with perm[...,maxi]. Step 3 above. perm = _swap_m_with_i(perm, m, maxi) # Step 4. row = batch_gather(matrix, perm[..., m:m + 1], axis=-2) row = batch_gather(row, perm[..., m + 1:]) # Step 5. prev_rows = pchol[..., :m, :] prev_rows_perm_m_onward = batch_gather(prev_rows, perm[..., m + 1:]) prev_rows_pivot_col = batch_gather(prev_rows, perm[..., m:m + 1]) row -= tf.reduce_sum(prev_rows_perm_m_onward * prev_rows_pivot_col, axis=-2)[..., tf.newaxis, :] # Step 6. pivot = tf.sqrt(maxval)[..., tf.newaxis, tf.newaxis] # Step 7. row = tf.concat([pivot, row / pivot], axis=-1) # TODO(b/130899118): Pad grad fails with int64 paddings. # Step 8. paddings = tf.concat([ tf.zeros([prefer_static.rank(pchol) - 1, 2], dtype=tf.int32), [[tf.cast(m, tf.int32), 0]] ], axis=0) diag_update = tf.pad(row**2, paddings=paddings)[..., 0, :] reverse_perm = _invert_permutation(perm) matrix_diag -= batch_gather(diag_update, reverse_perm) # Step 9. row = tf.pad(row, paddings=paddings) # TODO(bjp): Defer the reverse permutation all-at-once at the end? row = batch_gather(row, reverse_perm) pchol_shape = pchol.shape pchol = tf.concat([pchol[..., :m, :], row, pchol[..., m + 1:, :]], axis=-2) tensorshape_util.set_shape(pchol, pchol_shape) return m + 1, pchol, perm, matrix_diag
def _forward_log_det_jacobian(self, x): # For a discussion of this (non-obvious) result, see Note 7.2.2 (and the # sections leading up to it, for context) in # http://neutrino.aquaphoenix.com/ReactionDiffusion/SERC5chap7.pdf with tf.control_dependencies(self._assertions(x)): matrix_dim = tf.cast(tf.shape(x)[-1], dtype_util.base_dtype(x.dtype)) return -(matrix_dim + 1) * tf.reduce_sum( tf.math.log(tf.abs(tf.linalg.diag_part(x))), axis=-1)
def _reshape_part(part, dtype, event_shape): part = tf.cast(part, dtype) static_rank = tf.get_static_value( ps.rank_from_shape(event_shape)) if static_rank == 1: return part new_shape = ps.concat([ps.shape(part)[:-1], event_shape], axis=-1) return tf.reshape(part, ps.cast(new_shape, tf.int32))
def normal_conjugates_known_scale_posterior(prior, scale, s, n): """Posterior Normal distribution with conjugate prior on the mean. This model assumes that `n` observations (with sum `s`) come from a Normal with unknown mean `loc` (described by the Normal `prior`) and known variance `scale**2`. The "known scale posterior" is the distribution of the unknown `loc`. Accepts a prior Normal distribution object, having parameters `loc0` and `scale0`, as well as known `scale` values of the predictive distribution(s) (also assumed Normal), and statistical estimates `s` (the sum(s) of the observations) and `n` (the number(s) of observations). Returns a posterior (also Normal) distribution object, with parameters `(loc', scale'**2)`, where: ``` mu ~ N(mu', sigma'**2) sigma'**2 = 1/(1/sigma0**2 + n/sigma**2), mu' = (mu0/sigma0**2 + s/sigma**2) * sigma'**2. ``` Distribution parameters from `prior`, as well as `scale`, `s`, and `n`. will broadcast in the case of multidimensional sets of parameters. Args: prior: `Normal` object of type `dtype`: the prior distribution having parameters `(loc0, scale0)`. scale: tensor of type `dtype`, taking values `scale > 0`. The known stddev parameter(s). s: Tensor of type `dtype`. The sum(s) of observations. n: Tensor of type `int`. The number(s) of observations. Returns: A new Normal posterior distribution object for the unknown observation mean `loc`. Raises: TypeError: if dtype of `s` does not match `dtype`, or `prior` is not a Normal object. """ if not isinstance(prior, normal.Normal): raise TypeError("Expected prior to be an instance of type Normal") if s.dtype != prior.dtype: raise TypeError( "Observation sum s.dtype does not match prior dtype: %s vs. %s" % (s.dtype, prior.dtype)) n = tf.cast(n, prior.dtype) scale0_2 = tf.square(prior.scale) scale_2 = tf.square(scale) scalep_2 = 1.0 / (1 / scale0_2 + n / scale_2) return normal.Normal(loc=(prior.loc / scale0_2 + s / scale_2) * scalep_2, scale=tf.sqrt(scalep_2))
def _mean(self, probs=None): if probs is None: probs = self._categorical.probs_parameter() outcomes = self.outcomes if dtype_util.is_integer(outcomes.dtype): if self._validate_args: outcomes = dist_util.embed_check_integer_casting_closed( outcomes, target_dtype=probs.dtype) outcomes = tf.cast(outcomes, dtype=probs.dtype) return tf.tensordot(outcomes, probs, axes=[[0], [-1]])
def _inverse(self, y): # As specified in the Stan reference manual, the procedure is as follows: # N = y.shape[-1] # z_k = y_k / (1 - sum_{i=1 to k-1} y_i) # x_k = logit(z_k) - log(1 / (N - k)) offset = tf.math.log( tf.cast( tf.range(tf.shape(y)[-1] - 1, 0, delta=-1), dtype=dtype_util.base_dtype(y.dtype))) z = y / (1. - tf.math.cumsum(y, axis=-1, exclusive=True)) return tf.math.log(z[..., :-1]) - tf.math.log1p(-z[..., :-1]) + offset
def _sample_one_batch_member(args): logits, num_cat_samples = args[0], args[1] # [K], [] # x has shape [1, num_cat_samples = num_samples * num_trials] x = tf.random.categorical(logits[tf.newaxis, ...], num_cat_samples, seed=seed) x = tf.reshape(x, shape=[num_samples, -1]) # [num_samples, num_trials] x = tf.one_hot( x, depth=num_classes) # [num_samples, num_trials, num_classes] x = tf.reduce_sum(x, axis=-2) # [num_samples, num_classes] return tf.cast(x, dtype=dtype)
def _forward_log_det_jacobian(self, x): # is_constant_jacobian = True for this bijector, hence the # `log_det_jacobian` need only be specified for a single input, as this will # be tiled to match `event_ndims`. if self._is_only_identity_multiplier: # We don't pad in this case and instead let the fldj be applied # via broadcast. log_abs_diag = tf.math.log(tf.abs(self._scale)) event_size = tf.shape(x)[-1] event_size = tf.cast(event_size, dtype=log_abs_diag.dtype) return log_abs_diag * event_size return self.scale.log_abs_determinant()
def _inverse_log_det_jacobian(self, y): # Could also do: # ildj = tf.reduce_sum(y - tfp.math.softplus_inverse(y), # axis=event_dims) # but the following is more numerically stable. Ie, # Y = Log[1 + exp{X}] ==> X = Log[exp{Y} - 1] # ==> dX/dY = exp{Y} / (exp{Y} - 1) # = 1 / (1 - exp{-Y}), # which is the most stable for large Y > 0. For small Y, we use # 1 - exp{-Y} approx Y. if self.hinge_softness is not None: y = y / tf.cast(self.hinge_softness, y.dtype) return -tf.math.log(-tf.math.expm1(-y))
def _sample_n(self, n, seed=None): logits = self._logits_parameter_no_checks() logits_2d = tf.reshape(logits, [-1, self._num_categories(logits)]) sample_dtype = tf.int64 if dtype_util.size( self.dtype) > 4 else tf.int32 draws = tf.random.categorical(logits_2d, n, dtype=sample_dtype, seed=seed) draws = tf.cast(draws, self.dtype) return tf.reshape(tf.transpose(draws), shape=tf.concat( [[n], self._batch_shape_tensor(logits)], axis=0))
def _variance(self): probs = self._categorical.probs_parameter() outcomes = tf.broadcast_to(self.outcomes, shape=dist_util.prefer_static_shape(probs)) if dtype_util.is_integer(outcomes.dtype): if self._validate_args: outcomes = dist_util.embed_check_integer_casting_closed( outcomes, target_dtype=probs.dtype) outcomes = tf.cast(outcomes, dtype=probs.dtype) square_d = tf.math.squared_difference( outcomes, self._mean(probs)[..., tf.newaxis]) return tf.reduce_sum(probs * square_d, axis=-1)
def _forward_log_det_jacobian(self, x): # This code is similar to tf.math.log_softmax but different because we have # an implicit zero column to handle. I.e., instead of: # reduce_sum(logits - reduce_sum(exp(logits), dim)) # we must do: # log_normalization = 1 + reduce_sum(exp(logits)) # -log_normalization + reduce_sum(logits - log_normalization) n = prefer_static.shape(x)[-1] log_normalization = tf.math.softplus( tf.reduce_logsumexp(x, axis=-1, keepdims=True)) return tf.squeeze( (-log_normalization + tf.reduce_sum(x - log_normalization, axis=-1, keepdims=True)), axis=-1) + 0.5 * tf.math.log(tf.cast(n + 1, dtype=x.dtype))
def _cdf(self, x): # CDF(x) at positive integer x is the probability that the Zipf variable is # less than or equal to x; given by the formula: # CDF(x) = 1 - (zeta(power, x + 1) / Z) # For fractional x, the CDF is equal to the CDF at n = floor(x). # For x < 1, the CDF is zero. # If interpolate_nondiscrete is True, we return a continuous relaxation # which agrees with the CDF at integer points. power = tf.convert_to_tensor(self.power) x = tf.cast(x, power.dtype) safe_x = tf.maximum(x if self.interpolate_nondiscrete else tf.floor(x), 0.) cdf = 1. - (tf.math.zeta(power, safe_x + 1.) / tf.math.zeta(power, 1.)) return tf.where(x < 1., tf.zeros_like(cdf), cdf)
def _finish_prob_for_one_fiber(self, y, x, ildj, event_ndims, **distribution_kwargs): """Finish computation of prob on one element of the inverse image.""" x = self._maybe_rotate_dims(x, rotate_right=True) prob = self.distribution.prob(x, **distribution_kwargs) if self._is_maybe_event_override: prob = tf.reduce_prod(prob, axis=self._reduce_event_indices) prob = prob * tf.exp(tf.cast(ildj, prob.dtype)) if self._is_maybe_event_override and isinstance(event_ndims, int): tensorshape_util.set_shape( prob, tf.broadcast_static_shape( tensorshape_util.with_rank_at_least(y.shape, 1)[:-event_ndims], self.batch_shape)) return prob
def _forward(self, x): # As specified in the Stan reference manual, the procedure is as follows: # N = x.shape[-1] + 1 # z_k = sigmoid(x + log(1 / (N - k))) # y_1 = z_1 # y_k = (1 - sum_{i=1 to k-1} y_i) * z_k # y_N = 1 - sum_{i=1 to N-1} y_i # TODO(b/128857065): The numerics can possibly be improved here with a # log-space computation. offset = -tf.math.log( tf.cast( tf.range(tf.shape(x)[-1], 0, delta=-1), dtype=dtype_util.base_dtype(x.dtype))) z = tf.math.sigmoid(x + offset) y = z * tf.math.cumprod(1 - z, axis=-1, exclusive=True) return tf.concat([y, 1. - tf.reduce_sum(y, axis=-1, keepdims=True)], axis=-1)
def _log_normalization(self): """Computes the log-normalizer of the distribution.""" event_dim = tf.compat.dimension_value(self.event_shape[0]) if event_dim is None: raise ValueError('vMF _log_normalizer currently only supports ' 'statically known event shape') safe_conc = tf.where(self.concentration > 0, self.concentration, tf.ones_like(self.concentration)) safe_lognorm = ((event_dim / 2 - 1) * tf.math.log(safe_conc) - (event_dim / 2) * np.log(2 * np.pi) - tf.math.log( _bessel_ive(event_dim / 2 - 1, safe_conc)) - tf.abs(safe_conc)) log_nsphere_surface_area = ( np.log(2.) + (event_dim / 2) * np.log(np.pi) - tf.math.lgamma(tf.cast(event_dim / 2, self.dtype))) return tf.where(self.concentration > 0, -safe_lognorm, log_nsphere_surface_area)
def _prob(self, x): if self.validate_args: is_vector_check = assert_util.assert_rank_at_least(x, 1) right_vec_space_check = assert_util.assert_equal( self.event_shape_tensor(), tf.gather(tf.shape(x), tf.rank(x) - 1), message= "Argument 'x' not defined in the same space R^k as this distribution" ) with tf.control_dependencies([is_vector_check]): with tf.control_dependencies([right_vec_space_check]): x = tf.identity(x) loc = tf.convert_to_tensor(self.loc) return tf.cast(tf.reduce_all(tf.abs(x - loc) <= self._slack(loc), axis=-1), dtype=self.dtype)