def sign(x):
_check_arraylike('sign', x)
dtype = dtypes.dtype(x)
if dtypes.issubdtype(dtype, np.complexfloating):
re = lax.real(x)
return lax.complex(lax.sign(_where(re != 0, re, lax.imag(x))),
_constant_like(re, 0))
return lax.sign(x)
def floor_divide(x1, x2):
x1, x2 = _promote_args("floor_divide", x1, x2)
if onp.issubdtype(_dtype(x1), onp.integer):
quotient = lax.div(x1, x2)
select = logical_and(lax.sign(x1) != lax.sign(x2), lax.rem(x1, x2) != 0)
# TODO(mattjj): investigate why subtracting a scalar was causing promotion
return where(select, quotient - onp.array(1, _dtype(quotient)), quotient)
else:
return _float_divmod(x1, x2)[0]
def _float_divmod(x1, x2):
# see float_divmod in floatobject.c of CPython
mod = lax.rem(x1, x2)
div = lax.div(lax.sub(x1, mod), x2)
ind = lax.bitwise_and(mod != 0, lax.sign(x2) != lax.sign(mod))
mod = lax.select(ind, mod + x2, mod)
div = lax.select(ind, div - _constant_like(div, 1), div)
return lax.round(div), mod
def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
if b is not None:
a, b = jnp.broadcast_arrays(a, b)
dims = _reduction_dims(a, axis)
dimadd = lambda x: lax.expand_dims(x, dims)
amax = lax.reduce(a, _constant_like(a, -np.inf), lax.max, dims)
amax = lax.stop_gradient(
lax.select(lax.is_finite(amax), amax, lax.full_like(amax, 0)))
amax_singletons = dimadd(amax)
if b is None:
out = lax.add(
lax.log(
lax.reduce(lax.exp(lax.sub(a, amax_singletons)),
_constant_like(a, 0), lax.add, dims)), 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 = lax.reduce(lax.mul(lax.exp(lax.sub(a, amax_singletons)), b),
_constant_like(a, 0), lax.add, dims)
sign = lax.stop_gradient(lax.sign(sumexp))
out = lax.add(lax.log(lax.abs(sumexp)), amax)
if return_sign:
return (dimadd(out), dimadd(sign)) if keepdims else (out, sign)
if b is not None:
out = jnp.where(sign < 0, np.nan, out)
return dimadd(out) if keepdims else out
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 svd(a,
full_matrices: bool = True,
compute_uv: bool = True,
hermitian: bool = False):
a, = _promote_dtypes_inexact(jnp.asarray(a))
if hermitian:
w, v = lax_linalg.eigh(a)
s = lax.abs(v)
if compute_uv:
sign = lax.sign(v)
idxs = lax.broadcasted_iota(np.int64,
s.shape,
dimension=s.ndim - 1)
s, idxs, sign = lax.sort((s, idxs, sign), dimension=-1, num_keys=1)
s = lax.rev(s, dimensions=[s.ndim - 1])
idxs = lax.rev(idxs, dimensions=[s.ndim - 1])
sign = lax.rev(sign, dimensions=[s.ndim - 1])
u = jnp.take_along_axis(w, idxs[..., None, :], axis=-1)
vh = _H(u * sign[..., None, :])
return u, s, vh
else:
return lax.rev(lax.sort(s, dimension=-1), dimensions=[s.ndim - 1])
return lax_linalg.svd(a,
full_matrices=full_matrices,
compute_uv=compute_uv)
def floor_divide(x1, x2):
x1, x2 = _promote_args("floor_divide", x1, x2)
dtype = dtypes.dtype(x1)
if dtypes.issubdtype(dtype, np.integer):
quotient = lax.div(x1, x2)
select = logical_and(lax.sign(x1) != lax.sign(x2), lax.rem(x1, x2) != 0)
# TODO(mattjj): investigate why subtracting a scalar was causing promotion
return _where(select, quotient - 1, quotient)
elif dtypes.issubdtype(dtype, np.complexfloating):
x1r = lax.real(x1)
x1i = lax.imag(x1)
x2r = lax.real(x2)
x2i = lax.imag(x2)
which = lax.ge(lax.abs(x2r), lax.abs(x2i))
rat1 = _where(which, lax.full_like(x2i, 1), lax.div(x2r, x2i))
rat2 = _where(which, lax.div(x2i, x2r), _lax_const(x2i, 1))
out = lax.floor(lax.div(lax.add(lax.mul(x1r, rat1), lax.mul(x1i, rat2)),
lax.add(lax.mul(x2r, rat1), lax.mul(x2i, rat2))))
return lax.convert_element_type(out, dtype)
else:
return _float_divmod(x1, x2)[0]
def ldexp(x1, x2):
_check_arraylike("ldexp", x1, x2)
x1_dtype = dtypes.dtype(x1)
x2_dtype = dtypes.dtype(x2)
if (dtypes.issubdtype(x1_dtype, np.complexfloating)
or dtypes.issubdtype(x2_dtype, np.inexact)):
raise ValueError(
f"ldexp not supported for input types {(x1_dtype, x2_dtype)}")
x1, x2 = _promote_shapes("ldexp", x1, x2)
dtype = dtypes.canonicalize_dtype(dtypes._to_inexact_dtype(x1_dtype))
info = dtypes.finfo(dtype)
int_type = _INT_DTYPES[info.bits]
x1 = lax.convert_element_type(x1, dtype)
x2 = lax.convert_element_type(x2, int_type)
mask = (1 << info.nexp) - 1
bias = ((1 << info.nexp) - 1) >> 1
x, e = _normalize_float(x1)
x2 += e + ((x >> info.nmant) & mask) - bias
# find underflow/overflow before denormalization
underflow_cond = x2 < -(bias + info.nmant)
overflow_cond = x2 > bias
m = lax.full_like(x, 1, dtype=dtype)
# denormals
cond = x2 < -bias + 1
x2 = _where(cond, x2 + info.nmant, x2)
m = _where(cond, m / (1 << info.nmant), m)
x2 = lax.convert_element_type(x2, np.int32)
x &= ~(mask << info.nmant)
x |= ((lax.convert_element_type(x2, int_type) + bias) << info.nmant)
x = lax.convert_element_type(m, dtype) * lax.bitcast_convert_type(x, dtype)
# underflow
x = _where(underflow_cond, lax.full_like(x, 0, dtype=dtype), x)
# overflow
x = _where(overflow_cond, lax.sign(x1) * lax.full_like(x, np.inf), x)
# ldexp(x1, x2) = x1 for x1 = inf, -inf, nan, 0
return _where(isinf(x1) | isnan(x1) | (x1 == 0), x1, x)
def ldexp(x1, x2):
_check_arraylike("ldexp", x1, x2)
dtype = dtypes.canonicalize_dtype(_result_dtype(np.ldexp, x1, x2))
x1, x2 = _promote_shapes("ldexp", x1, x2)
x1 = lax.convert_element_type(x1, dtype)
info = dtypes.finfo(dtype)
mask = (1 << info.nexp) - 1
bias = ((1 << info.nexp) - 1) >> 1
int_type = _INT_DTYPES[info.bits]
x, e = _normalize_float(x1)
x2 += e + ((x >> info.nmant) & mask) - bias
# find underflow/overflow before denormalization
underflow_cond = x2 < -(bias + info.nmant)
overflow_cond = x2 > bias
m = lax.full_like(x, 1, dtype=dtype)
# denormals
cond = x2 < -bias + 1
x2 = _where(cond, x2 + info.nmant, x2)
m = _where(cond, m / (1 << info.nmant), m)
x2 = lax.convert_element_type(x2, np.int32)
x &= ~(mask << info.nmant)
x |= ((lax.convert_element_type(x2, int_type) + bias) << info.nmant)
x = lax.convert_element_type(m, dtype) * lax.bitcast_convert_type(x, dtype)
# underflow
x = _where(underflow_cond, lax.full_like(x, 0, dtype=dtype), x)
# overflow
x = _where(overflow_cond, lax.sign(x1) * lax.full_like(x, np.inf), x)
# ldexp(x1, x2) = x1 for x1 = inf, -inf, nan, 0
return _where(isinf(x1) | isnan(x1) | (x1 == 0), x1, x)
def _abs2(x):
x, xx = x
sign = jnp.where(lax.sign(x) == lax.sign(xx), 1, -1)
return (lax.abs(x), sign * lax.abs(xx))
def cdf(x, p):
x, p = _promote_args_inexact("gennorm.cdf", x, p)
return .5 * (1 + lax.sign(x) * lax.igamma(1 / p, lax.abs(x)**p))
def correct_rank_deficiency(u_out):
u_out, r = lax.linalg.qr(u_out, full_matrices=False)
u_out = u_out @ jnp.diag(lax.sign(jnp.diag(r)))
return u_out
