def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
if b is not None:
a, b = _promote_args_inexact("logsumexp", a, b)
a = jnp.where(b != 0, a, -jnp.inf)
pos_dims, dims = _reduction_dims(a, axis)
amax = jnp.max(a, axis=dims, keepdims=keepdims)
amax = lax.stop_gradient(
lax.select(lax.is_finite(amax), amax, lax.full_like(amax, 0)))
amax_with_dims = amax if keepdims else lax.expand_dims(amax, pos_dims)
if b is None:
out = lax.add(
lax.log(
jnp.sum(lax.exp(lax.sub(a, amax_with_dims)),
axis=dims,
keepdims=keepdims)), amax)
sign = jnp.where(jnp.isnan(out), np.nan, 1.0).astype(out.dtype)
sign = jnp.where(out == -np.inf, 0.0, sign)
else:
sumexp = jnp.sum(lax.mul(lax.exp(lax.sub(a, amax_with_dims)), b),
axis=dims,
keepdims=keepdims)
sign = lax.stop_gradient(lax.sign(sumexp))
out = lax.add(lax.log(lax.abs(sumexp)), amax)
if return_sign:
return (out, sign)
if b is not None:
out = jnp.where(sign < 0, np.nan, out)
return out
def logpdf(x, alpha):
args = (np.ones((0, ), lax.dtype(x)), np.ones((1, ), lax.dtype(alpha)))
to_dtype = lax.dtype(osp_stats.dirichlet.logpdf(*args))
x, alpha = [lax.convert_element_type(arg, to_dtype) for arg in (x, alpha)]
one = jnp._constant_like(x, 1)
normalize_term = jnp.sum(gammaln(alpha), axis=-1) - gammaln(
jnp.sum(alpha, axis=-1))
log_probs = lax.sub(jnp.sum(xlogy(lax.sub(alpha, one), x), axis=-1),
normalize_term)
return jnp.where(_is_simplex(x), log_probs, -jnp.inf)
def segment_sum(data,
segment_ids,
num_segments=None,
indices_are_sorted=False,
unique_indices=False,
bucket_size=None): # TODO(zhangqiaorjc): use non-None default.
"""Computes the sum within segments of an array.
Similar to TensorFlow's segment_sum:
https://www.tensorflow.org/api_docs/python/tf/math/segment_sum
Args:
data: an array with the values to be summed.
segment_ids: an array with integer dtype that indicates the segments of
`data` (along its leading axis) to be summed. Values can be repeated and
need not be sorted. Values outside of the range [0, num_segments) are
dropped and do not contribute to the sum.
num_segments: optional, an int with nonnegative value indicating the number
of segments. The default is set to be the minimum number of segments that
would support all indices in ``segment_ids``, calculated as
``max(segment_ids) + 1``.
Since `num_segments` determines the size of the output, a static value
must be provided to use ``segment_sum`` in a ``jit``-compiled function.
indices_are_sorted: whether ``segment_ids`` is known to be sorted.
unique_indices: whether `segment_ids` is known to be free of duplicates.
bucket_size: size of bucket to group indices into. ``segment_sum`` is
performed on each bucket separately to improve numerical stability of
addition. Default ``None`` means no bucketing.
Returns:
An array with shape :code:`(num_segments,) + data.shape[1:]` representing the
segment sums.
"""
if num_segments is None:
num_segments = jnp.max(segment_ids) + 1
num_segments = int(num_segments)
out = jnp.zeros((num_segments, ) + data.shape[1:], dtype=data.dtype)
num_buckets = 1 if bucket_size is None \
else util.ceil_of_ratio(segment_ids.size, bucket_size)
if num_buckets == 1:
return _scatter_update(out,
segment_ids,
data,
lax.scatter_add,
indices_are_sorted,
unique_indices,
normalize_indices=False)
# Bucketize indices and perform segment_sum on each bucket to improve
# numerical stability.
outs = []
for sub_data, sub_segment_ids in zip(
jnp.array_split(data, num_buckets),
jnp.array_split(segment_ids, num_buckets)):
outs.append(
segment_sum(sub_data, sub_segment_ids, num_segments,
indices_are_sorted, unique_indices))
return jnp.sum(jnp.stack(outs), axis=0)
def zeta(x, q=None):
assert q is not None, "Riemann zeta function is not implemented yet."
# Reference: Johansson, Fredrik.
# "Rigorous high-precision computation of the Hurwitz zeta function and its derivatives."
# Numerical Algorithms 69.2 (2015): 253-270.
# https://arxiv.org/abs/1309.2877 - formula (5)
# here we keep the same notation as in reference
s, a = _promote_args_inexact("zeta", x, q)
dtype = lax.dtype(a).type
s_, a_ = jnp.expand_dims(s, -1), jnp.expand_dims(a, -1)
# precision ~ N, M
N = M = dtype(8) if lax.dtype(a) == jnp.float32 else dtype(16)
assert M <= len(_BERNOULLI_COEFS)
k = jnp.expand_dims(np.arange(N, dtype=N.dtype), tuple(range(a.ndim)))
S = jnp.sum((a_ + k)**-s_, -1)
I = lax.div((a + N)**(dtype(1) - s), s - dtype(1))
T0 = (a + N)**-s
m = jnp.expand_dims(np.arange(2 * M, dtype=M.dtype), tuple(range(s.ndim)))
s_over_a = (s_ + m) / (a_ + N)
T1 = jnp.cumprod(s_over_a, -1)[..., ::2]
T1 = jnp.clip(T1, a_max=jnp.finfo(dtype).max)
coefs = np.expand_dims(
np.array(_BERNOULLI_COEFS[:T1.shape[-1]], dtype=dtype),
tuple(range(a.ndim)))
T1 = T1 / coefs
T = T0 * (dtype(0.5) + T1.sum(-1))
return S + I + T
def _all_gather_transpose_rule(cts, x, *, all_gather_dimension, axis_name,
axis_index_groups, axis_size):
# TODO(cjfj): Add reduce-scatter op to XLA?
concat_axis = 0
return (lax_numpy.sum(all_to_all(cts,
axis_name=axis_name,
split_axis=all_gather_dimension,
concat_axis=concat_axis,
axis_index_groups=axis_index_groups),
axis=concat_axis), )
def matrix_rank(M, tol=None):
M = _promote_arg_dtypes(jnp.asarray(M))
if M.ndim > 2:
raise TypeError("array should have 2 or fewer dimensions")
if M.ndim < 2:
return jnp.any(M != 0).astype(jnp.int32)
S = svd(M, full_matrices=False, compute_uv=False)
if tol is None:
tol = S.max() * np.max(M.shape) * jnp.finfo(S.dtype).eps
return jnp.sum(S > tol)
def logpdf(x, alpha):
x, alpha = _promote_dtypes_inexact(x, alpha)
if alpha.ndim != 1:
raise ValueError(
f"`alpha` must be one-dimensional; got alpha.shape={alpha.shape}"
)
if x.shape[0] not in (alpha.shape[0], alpha.shape[0] - 1):
raise ValueError(
"`x` must have either the same number of entries as `alpha` "
f"or one entry fewer; got x.shape={x.shape}, alpha.shape={alpha.shape}"
)
one = lax._const(x, 1)
if x.shape[0] != alpha.shape[0]:
x = jnp.concatenate([x, lax.sub(one, x.sum(0, keepdims=True))], axis=0)
normalize_term = jnp.sum(gammaln(alpha)) - gammaln(jnp.sum(alpha))
if x.ndim > 1:
alpha = lax.broadcast_in_dim(alpha, alpha.shape + (1,) * (x.ndim - 1), (0,))
log_probs = lax.sub(jnp.sum(xlogy(lax.sub(alpha, one), x), axis=0), normalize_term)
return jnp.where(_is_simplex(x), log_probs, -jnp.inf)
def multigammaln(a, d):
d = core.concrete_or_error(int, d, "d argument of multigammaln")
a, d = _promote_args_inexact("multigammaln", a, d)
constant = lax.mul(lax.mul(lax.mul(_constant_like(a, 0.25), d),
lax.sub(d, _constant_like(a, 1))),
lax.log(_constant_like(a, np.pi)))
res = jnp.sum(gammaln(jnp.expand_dims(a, axis=-1) -
lax.div(jnp.arange(d), _constant_like(a, 2))),
axis=-1)
return res + constant
def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
if b is not None:
a, b = _promote_args_inexact("logsumexp", a, b)
a = jnp.where(b != 0, a, -jnp.inf)
else:
a, = _promote_args_inexact("logsumexp", a)
pos_dims, dims = _reduction_dims(a, axis)
amax = jnp.max(a, axis=dims, keepdims=keepdims)
amax = lax.stop_gradient(
lax.select(jnp.isfinite(amax), amax, lax.full_like(amax, 0)))
amax_with_dims = amax if keepdims else lax.expand_dims(amax, pos_dims)
# fast path if the result cannot be negative.
if b is None and not np.issubdtype(a.dtype, np.complexfloating):
out = lax.add(
lax.log(
jnp.sum(lax.exp(lax.sub(a, amax_with_dims)),
axis=dims,
keepdims=keepdims)), amax)
sign = jnp.where(jnp.isnan(out), out, 1.0)
sign = jnp.where(jnp.isneginf(out), 0.0, sign).astype(out.dtype)
else:
expsub = lax.exp(lax.sub(a, amax_with_dims))
if b is not None:
expsub = lax.mul(expsub, b)
sumexp = jnp.sum(expsub, axis=dims, keepdims=keepdims)
sign = lax.stop_gradient(jnp.sign(sumexp))
if np.issubdtype(sumexp.dtype, np.complexfloating):
if return_sign:
sumexp = sign * sumexp
out = lax.add(lax.log(sumexp), amax)
else:
out = lax.add(lax.log(lax.abs(sumexp)), amax)
if return_sign:
return (out, sign)
if b is not None:
if not np.issubdtype(out.dtype, np.complexfloating):
# Use jnp.array(nan) to avoid false positives in debug_nans
# (see https://github.com/google/jax/issues/7634)
out = jnp.where(sign < 0, jnp.array(np.nan, dtype=out.dtype), out)
return out
def multigammaln(a, d):
a, = _promote_args_inexact("multigammaln", a)
d = lax.convert_element_type(d, lax.dtype(a))
constant = lax.mul(
lax.mul(lax.mul(_constant_like(a, 0.25), d),
lax.sub(d, _constant_like(a, 1))),
lax.log(_constant_like(a, np.pi)))
res = jnp.sum(gammaln(
jnp.expand_dims(a, axis=-1) -
lax.div(jnp.arange(d), _constant_like(a, 2))),
axis=-1)
return res + constant
def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
if b is not None:
a, b = _promote_args_inexact("logsumexp", a, b)
a = jnp.where(b != 0, a, -jnp.inf)
else:
a, = _promote_args_inexact("logsumexp", a)
pos_dims, dims = _reduction_dims(a, axis)
amax = jnp.max(a, axis=dims, keepdims=keepdims)
amax = lax.stop_gradient(
lax.select(jnp.isfinite(amax), amax, lax.full_like(amax, 0)))
amax_with_dims = amax if keepdims else lax.expand_dims(amax, pos_dims)
# fast path if the result cannot be negative.
if b is None and not np.issubdtype(a.dtype, np.complexfloating):
out = lax.add(
lax.log(
jnp.sum(lax.exp(lax.sub(a, amax_with_dims)),
axis=dims,
keepdims=keepdims)), amax)
sign = jnp.where(jnp.isnan(out), np.nan, 1.0).astype(out.dtype)
sign = jnp.where(out == -np.inf, 0.0, sign)
else:
expsub = lax.exp(lax.sub(a, amax_with_dims))
if b is not None:
expsub = lax.mul(expsub, b)
sumexp = jnp.sum(expsub, axis=dims, keepdims=keepdims)
sign = lax.stop_gradient(jnp.sign(sumexp))
if np.issubdtype(sumexp.dtype, np.complexfloating):
if return_sign:
sumexp = sign * sumexp
out = lax.add(lax.log(sumexp), amax)
else:
out = lax.add(lax.log(lax.abs(sumexp)), amax)
if return_sign:
return (out, sign)
if b is not None:
if not np.issubdtype(out.dtype, np.complexfloating):
out = jnp.where(sign < 0, np.nan, out)
return out
def eig_jvp_rule(primals, tangents, *, compute_left_eigenvectors,
compute_right_eigenvectors):
if compute_left_eigenvectors or compute_right_eigenvectors:
raise NotImplementedError(
'The derivatives of eigenvectors are not implemented, only '
'eigenvalues. See '
'https://github.com/google/jax/issues/2748 for discussion.')
# Formula for derivative of eigenvalues w.r.t. a is eqn 4.60 in
# https://arxiv.org/abs/1701.00392
a, = primals
da, = tangents
l, v = eig(a, compute_left_eigenvectors=False)
return [l], [jnp.sum(_solve(v, da.astype(v.dtype)) * _T(v), -1)]
def multigammaln(a, d):
d = core.concrete_or_error(int, d, "d argument of multigammaln")
a, d_ = _promote_args_inexact("multigammaln", a, d)
constant = lax.mul(
lax.mul(lax.mul(lax._const(a, 0.25), d_),
lax.sub(d_, lax._const(a, 1))), lax.log(lax._const(a, np.pi)))
b = lax.div(jnp.arange(d, dtype=d_.dtype), lax._const(a, 2))
res = jnp.sum(gammaln(
jnp.expand_dims(a, axis=-1) -
jnp.expand_dims(b, axis=tuple(range(a.ndim)))),
axis=-1)
return res + constant
def _slogdet_lu(a):
dtype = lax.dtype(a)
lu, pivot, _ = lax_linalg.lu(a)
diag = jnp.diagonal(lu, axis1=-2, axis2=-1)
is_zero = jnp.any(diag == jnp.array(0, dtype=dtype), axis=-1)
iota = lax.expand_dims(jnp.arange(a.shape[-1]), range(pivot.ndim - 1))
parity = jnp.count_nonzero(pivot != iota, axis=-1)
if jnp.iscomplexobj(a):
sign = jnp.prod(diag / jnp.abs(diag), axis=-1)
else:
sign = jnp.array(1, dtype=dtype)
parity = parity + jnp.count_nonzero(diag < 0, axis=-1)
sign = jnp.where(is_zero, jnp.array(0, dtype=dtype),
sign * jnp.array(-2 * (parity % 2) + 1, dtype=dtype))
logdet = jnp.where(is_zero, jnp.array(-jnp.inf, dtype=dtype),
jnp.sum(jnp.log(jnp.abs(diag)), axis=-1))
return sign, jnp.real(logdet)
def _median_bias(n):
"""
Returns the bias of the median of a set of periodograms relative to
the mean. See Appendix B from [1]_ for details.
Args:
n : int
Numbers of periodograms being averaged.
Returns:
bias : float
Calculated bias.
References:
.. [1] B. Allen, W.G. Anderson, P.R. Brady, D.A. Brown, J.D.E. Creighton.
"FINDCHIRP: an algorithm for detection of gravitational waves from
inspiraling compact binaries", Physical Review D 85, 2012,
:arxiv:`gr-qc/0509116`
"""
ii_2 = jnp.arange(2., n, 2)
return 1 + jnp.sum(1. / (ii_2 + 1) - 1. / ii_2)
def slogdet(a):
a = _promote_arg_dtypes(jnp.asarray(a))
dtype = lax.dtype(a)
a_shape = jnp.shape(a)
if len(a_shape) < 2 or a_shape[-1] != a_shape[-2]:
msg = "Argument to slogdet() must have shape [..., n, n], got {}"
raise ValueError(msg.format(a_shape))
lu, pivot, _ = lax_linalg.lu(a)
diag = jnp.diagonal(lu, axis1=-2, axis2=-1)
is_zero = jnp.any(diag == jnp.array(0, dtype=dtype), axis=-1)
parity = jnp.count_nonzero(pivot != jnp.arange(a_shape[-1]), axis=-1)
if jnp.iscomplexobj(a):
sign = jnp.prod(diag / jnp.abs(diag), axis=-1)
else:
sign = jnp.array(1, dtype=dtype)
parity = parity + jnp.count_nonzero(diag < 0, axis=-1)
sign = jnp.where(is_zero, jnp.array(0, dtype=dtype),
sign * jnp.array(-2 * (parity % 2) + 1, dtype=dtype))
logdet = jnp.where(is_zero, jnp.array(-jnp.inf, dtype=dtype),
jnp.sum(jnp.log(jnp.abs(diag)), axis=-1))
return sign, jnp.real(logdet)
def _is_simplex(x): x_sum = jnp.sum(x, axis=0) return jnp.all(x > 0, axis=0) & (abs(x_sum - 1) < 1E-6)
def norm(x,
ord=None,
axis: Union[None, Tuple[int, ...], int] = None,
keepdims=False):
x = _promote_arg_dtypes(jnp.asarray(x))
x_shape = jnp.shape(x)
ndim = len(x_shape)
if axis is None:
# NumPy has an undocumented behavior that admits arbitrary rank inputs if
# `ord` is None: https://github.com/numpy/numpy/issues/14215
if ord is None:
return jnp.sqrt(
jnp.sum(jnp.real(x * jnp.conj(x)), keepdims=keepdims))
axis = tuple(range(ndim))
elif isinstance(axis, tuple):
axis = tuple(canonicalize_axis(x, ndim) for x in axis)
else:
axis = (canonicalize_axis(axis, ndim), )
num_axes = len(axis)
if num_axes == 1:
if ord is None or ord == 2:
return jnp.sqrt(
jnp.sum(jnp.real(x * jnp.conj(x)),
axis=axis,
keepdims=keepdims))
elif ord == jnp.inf:
return jnp.amax(jnp.abs(x), axis=axis, keepdims=keepdims)
elif ord == -jnp.inf:
return jnp.amin(jnp.abs(x), axis=axis, keepdims=keepdims)
elif ord == 0:
return jnp.sum(x != 0,
dtype=jnp.finfo(lax.dtype(x)).dtype,
axis=axis,
keepdims=keepdims)
elif ord == 1:
# Numpy has a special case for ord == 1 as an optimization. We don't
# really need the optimization (XLA could do it for us), but the Numpy
# code has slightly different type promotion semantics, so we need a
# special case too.
return jnp.sum(jnp.abs(x), axis=axis, keepdims=keepdims)
else:
abs_x = jnp.abs(x)
ord = lax._const(abs_x, ord)
out = jnp.sum(abs_x**ord, axis=axis, keepdims=keepdims)
return jnp.power(out, 1. / ord)
elif num_axes == 2:
row_axis, col_axis = cast(Tuple[int, ...], axis)
if ord is None or ord in ('f', 'fro'):
return jnp.sqrt(
jnp.sum(jnp.real(x * jnp.conj(x)),
axis=axis,
keepdims=keepdims))
elif ord == 1:
if not keepdims and col_axis > row_axis:
col_axis -= 1
return jnp.amax(jnp.sum(jnp.abs(x),
axis=row_axis,
keepdims=keepdims),
axis=col_axis,
keepdims=keepdims)
elif ord == -1:
if not keepdims and col_axis > row_axis:
col_axis -= 1
return jnp.amin(jnp.sum(jnp.abs(x),
axis=row_axis,
keepdims=keepdims),
axis=col_axis,
keepdims=keepdims)
elif ord == jnp.inf:
if not keepdims and row_axis > col_axis:
row_axis -= 1
return jnp.amax(jnp.sum(jnp.abs(x),
axis=col_axis,
keepdims=keepdims),
axis=row_axis,
keepdims=keepdims)
elif ord == -jnp.inf:
if not keepdims and row_axis > col_axis:
row_axis -= 1
return jnp.amin(jnp.sum(jnp.abs(x),
axis=col_axis,
keepdims=keepdims),
axis=row_axis,
keepdims=keepdims)
elif ord in ('nuc', 2, -2):
x = jnp.moveaxis(x, axis, (-2, -1))
if ord == 2:
reducer = jnp.amax
elif ord == -2:
reducer = jnp.amin
else:
reducer = jnp.sum
y = reducer(svd(x, compute_uv=False), axis=-1)
if keepdims:
result_shape = list(x_shape)
result_shape[axis[0]] = 1
result_shape[axis[1]] = 1
y = jnp.reshape(y, result_shape)
return y
else:
raise ValueError("Invalid order '{}' for matrix norm.".format(ord))
else:
raise ValueError(
"Invalid axis values ({}) for jnp.linalg.norm.".format(axis))
def _is_simplex(x):
x_sum = jnp.sum(x, axis=-1)
return jnp.all(x > 0, axis=-1) & (x_sum <= 1) & (x_sum > 1 - 1e-6)
