def _mul_sparse(spenv, *argspecs): X, Y = argspecs if X.is_sparse() and Y.is_sparse(): if X.shape != Y.shape: raise NotImplementedError( "Multiplication between sparse matrices of different shapes.") if X.indices_ref == Y.indices_ref: out_data = lax.mul(X.data(spenv), Y.data(spenv)) out_argspec = ArgSpec(X.shape, spenv.push(out_data), X.indices_ref) elif X.indices(spenv).ndim != Y.indices(spenv).ndim or X.data( spenv).ndim != Y.data(spenv).ndim: raise NotImplementedError( "Multiplication between sparse matrices with different batch/dense dimensions." ) else: raise NotImplementedError( "Multiplication between sparse matrices with different sparsity patterns." ) else: if Y.is_sparse(): X, Y = Y, X Ydata = Y.data(spenv) if Ydata.ndim == 0: out_data = lax.mul(X.data(spenv), Ydata) elif Ydata.shape == X.shape: out_data = lax.mul(X.data(spenv), sparse.bcoo_extract(X.indices(spenv), Ydata)) else: raise NotImplementedError( "Multiplication between sparse and dense matrices of different shape." ) out_argspec = ArgSpec(X.shape, spenv.push(out_data), X.indices_ref) return (out_argspec, )
def logpdf(x, loc=0, scale=1): x, loc, scale = _promote_args_inexact("cauchy.logpdf", x, loc, scale) pi = _constant_like(x, np.pi) scaled_x = lax.div(lax.sub(x, loc), scale) normalize_term = lax.log(lax.mul(pi, scale)) return lax.neg( lax.add(normalize_term, lax.log1p(lax.mul(scaled_x, scaled_x))))
def cfun(x): return lax.cond( lax.lt(x, 2), x, lambda x: lax.mul(2, x), x, lambda x: lax.cond(lax.lt(x, 5), x, lambda x: lax.mul(3, x), 4, lambda y: lax.mul(y, x)))
def inner_cond(x): return lax.cond( lax.lt(x, 5), x, lambda x: lax.mul(3, x), 4, lambda y: lax.mul(y, x), )
def cdf(x, loc=0, scale=1): x, loc, scale = _promote_args_inexact("laplace.cdf", x, loc, scale) half = _constant_like(x, 0.5) one = _constant_like(x, 1) zero = _constant_like(x, 0) diff = lax.div(lax.sub(x, loc), scale) return lax.select(lax.le(diff, zero), lax.mul(half, lax.exp(diff)), lax.sub(one, lax.mul(half, lax.exp(lax.neg(diff)))))
def fun(x): if x < 2: return lax.mul(2, x) else: if x < 5: return lax.mul(3, x) else: return lax.mul(4, x)
def _logaddexp2_jvp(primals, tangents): x1, x2 = primals t1, t2 = tangents x1, x2, t1, t2 = _promote_args_inexact("logaddexp2_jvp", x1, x2, t1, t2) primal_out = logaddexp2(x1, x2) tangent_out = lax.add(lax.mul(t1, exp2(lax.sub(_replace_inf(x1), _replace_inf(primal_out)))), lax.mul(t2, exp2(lax.sub(_replace_inf(x2), _replace_inf(primal_out))))) return primal_out, tangent_out
def fun(x): res = 0 if x < 2: res = lax.mul(2, x) else: if x < 5: res = lax.mul(3, x) else: res = lax.mul(4, x) return res
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 logpdf(x, a, b, loc=0, scale=1): x, a, b, loc, scale = _promote_args_inexact("beta.logpdf", x, a, b, loc, scale) one = _constant_like(x, 1) shape_term = lax.neg(betaln(a, b)) y = lax.div(lax.sub(x, loc), scale) log_linear_term = lax.add(lax.mul(lax.sub(a, one), lax.log(y)), lax.mul(lax.sub(b, one), lax.log1p(lax.neg(y)))) log_probs = lax.sub(lax.add(shape_term, log_linear_term), lax.log(scale)) return where(logical_or(lax.gt(x, lax.add(loc, scale)), lax.lt(x, loc)), -inf, log_probs)
def logpdf(x, df, loc=0, scale=1): x, df, loc, scale = _promote_args_inexact("t.logpdf", x, df, loc, scale) two = _lax_const(x, 2) scaled_x = lax.div(lax.sub(x, loc), scale) df_over_two = lax.div(df, two) df_plus_one_over_two = lax.add(df_over_two, _lax_const(x, 0.5)) normalize_term_const = lax.mul(lax.mul(scale, scale), _lax_const(x, np.pi)) normalize_term_tmp = lax.div(lax.log(lax.mul(normalize_term_const, df)), two) normalize_term = lax.sub(lax.add(lax.lgamma(df_over_two), normalize_term_tmp), lax.lgamma(df_plus_one_over_two)) quadratic = lax.div(lax.mul(scaled_x, scaled_x), df) return lax.neg(lax.add(normalize_term, lax.mul(df_plus_one_over_two, lax.log1p(quadratic))))
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 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 logpdf(x, df, loc=0, scale=1): x, df, loc, scale = _promote_args_inexact("chi2.logpdf", x, df, loc, scale) one = _constant_like(x, 1) two = _constant_like(x, 2) y = lax.div(lax.sub(x, loc), scale) df_on_two = lax.div(df, two) kernel = lax.sub(lax.mul(lax.sub(df_on_two, one), lax.log(y)), lax.div(y,two)) nrml_cnst = lax.neg(lax.add(lax.lgamma(df_on_two),lax.div(lax.mul(lax.log(two), df),two))) log_probs = lax.add(lax.sub(nrml_cnst, lax.log(scale)), kernel) return where(lax.lt(x, loc), -inf, log_probs)
def _mul_sparse(spenv, *spvalues): X, Y = spvalues if X.is_sparse() and Y.is_sparse(): if X.indices_ref == Y.indices_ref and X.unique_indices: if config.jax_enable_checks: assert X.indices_sorted == Y.indices_sorted assert X.unique_indices == Y.unique_indices out_data = lax.mul(spenv.data(X), spenv.data(Y)) out_spvalue = spenv.sparse(X.shape, out_data, indices_ref=X.indices_ref, indices_sorted=X.indices_sorted, unique_indices=True) else: X_promoted, Y_promoted = spvalues_to_arrays(spenv, spvalues) mat = bcoo_multiply_sparse(X_promoted, Y_promoted) out_spvalue = spenv.sparse(mat.shape, mat.data, mat.indices) else: if Y.is_sparse(): X, Y = Y, X X_promoted = spvalues_to_arrays(spenv, X) out_data = bcoo_multiply_dense(X_promoted, spenv.data(Y)) out_spvalue = spenv.sparse(X.shape, out_data, indices_ref=X.indices_ref, indices_sorted=X.indices_sorted, unique_indices=X.unique_indices) return (out_spvalue, )
def logpdf(x, loc=0, scale=1): x, loc, scale = _promote_args_inexact("norm.logpdf", x, loc, scale) two = _constant_like(x, 2) scale_sqrd = lax.pow(scale, two) log_normalizer = lax.log(lax.mul(_constant_like(x, 2 * np.pi), scale_sqrd)) quadratic = lax.div(lax.pow(lax.sub(x, loc), two), scale_sqrd) return lax.div(lax.neg(lax.add(log_normalizer, quadratic)), two)
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 nanvar(a, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype=None, out=None, ddof=0, keepdims=False, where=None): _check_arraylike("nanvar", a) lax_internal._check_user_dtype_supported(dtype, "nanvar") if out is not None: raise NotImplementedError( "The 'out' argument to jnp.nanvar is not supported.") a_dtype, dtype = _var_promote_types(dtypes.dtype(a), dtype) a_mean = nanmean(a, axis, dtype=a_dtype, keepdims=True, where=where) centered = _where(lax_internal._isnan(a), 0, a - a_mean) # double-where trick for gradients. if dtypes.issubdtype(centered.dtype, np.complexfloating): centered = lax.real(lax.mul(centered, lax.conj(centered))) else: centered = lax.square(centered) normalizer = sum(lax_internal.bitwise_not(lax_internal._isnan(a)), axis=axis, keepdims=keepdims, where=where) normalizer = normalizer - ddof normalizer_mask = lax.le(normalizer, 0) result = sum(centered, axis, keepdims=keepdims, where=where) result = _where(normalizer_mask, np.nan, result) divisor = _where(normalizer_mask, 1, normalizer) out = lax.div(result, lax.convert_element_type(divisor, result.dtype)) return lax.convert_element_type(out, dtype)
def var(a, axis=None, keepdims=False, ddof=0): if ddof != 0: raise NotImplementedError("Only implemented for ddof=0.") centered = subtract(a, mean(a, axis, keepdims=True)) if iscomplexobj(centered): centered = lax.abs(centered) return mean(lax.mul(centered, centered), axis, keepdims=keepdims)
def xlog1py(x, y): x, y = _promote_args_inexact("xlog1py", x, y) x_ok = x != 0. safe_x = jnp.where(x_ok, x, 1.) safe_y = jnp.where(x_ok, y, 1.) return jnp.where(x_ok, lax.mul(safe_x, lax.log1p(safe_y)), jnp.zeros_like(x))
def _mul_sparse(spenv, *argspecs): X, Y = argspecs if X.is_sparse() and Y.is_sparse(): if X.indices_ref == Y.indices_ref: # TODO(jakevdp): this is inaccurate if there are duplicate indices out_data = lax.mul(X.data(spenv), Y.data(spenv)) out_argspec = ArgSpec(X.shape, spenv.push(out_data), X.indices_ref) else: data, indices, shape = bcoo_multiply_sparse(X.data(spenv), X.indices(spenv), Y.data(spenv), Y.indices(spenv), lhs_shape=X.shape, rhs_shape=Y.shape) out_argspec = ArgSpec(shape, spenv.push(data), spenv.push(indices)) else: if Y.is_sparse(): X, Y = Y, X out_data = bcoo_multiply_dense(X.data(spenv), X.indices(spenv), Y.data(spenv), shape=X.shape) out_argspec = ArgSpec(X.shape, spenv.push(out_data), X.indices_ref) return (out_argspec, )
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 exponential(tensor, dtype, exp_hparams): """Calculates an exponential approximation based on exp hyper params.""" # If low_bound defined, it clips x-M. if exp_hparams.low_bound != 0: tensor = jnp.clip(tensor, exp_hparams.low_bound, 0.) # TODO(luispazos) Use standard calls to top level jnp functions. # pylint: disable=protected-access def make_constant(c): return lax_numpy._constant_like(tensor, c).astype(dtype) # If clip_and_subtract, replace exp(clip(x-M,low_bound)) term with # exp(clip(x-M,low_bound))-exp(low_bound).' if exp_hparams.clip_and_subtract: tensor = lax.sub(tensor, make_constant(onp.exp(exp_hparams.low_bound))) # If linear_gradient: use this gradient as linear approximation of # exponential. if exp_hparams.linear_gradient is not None and exp_hparams.linear_gradient != 0: # Want: max(0, a*x+b) such that a*x+b goes through (0, 1). # # This comes out to: max(0, a*x+1), for arbitrary a>0. one = jnp.full(tensor.shape, 1.).astype(dtype) gradient = jnp.full(tensor.shape, exp_hparams.linear_gradient).astype(dtype) approx_exp = jnp.clip(lax.add(lax.mul(tensor, gradient), one), 0, 1) else: approx_exp = lax.exp(tensor) return approx_exp
def sinc(x): _check_arraylike("sinc", x) x, = _promote_dtypes_inexact(x) eq_zero = lax.eq(x, _lax_const(x, 0)) pi_x = lax.mul(_lax_const(x, np.pi), x) safe_pi_x = _where(eq_zero, _lax_const(x, 1), pi_x) return _where(eq_zero, _sinc_maclaurin(0, pi_x), lax.div(lax.sin(safe_pi_x), safe_pi_x))
def logpdf(x, a, loc=0, scale=1): x, a, loc, scale = _promote_args_inexact("gamma.logpdf", x, a, loc, scale) one = _constant_like(x, 1) y = lax.div(lax.sub(x, loc), scale) log_linear_term = lax.sub(lax.mul(lax.sub(a, one), lax.log(y)), y) shape_terms = lax.add(gammaln(a), lax.log(scale)) log_probs = lax.sub(log_linear_term, shape_terms) return where(lax.lt(x, loc), -inf, log_probs)
def logpdf(x, b, loc=0, scale=1): x, b, loc, scale = _promote_args_inexact("pareto.logpdf", x, b, loc, scale) one = _constant_like(x, 1) scaled_x = lax.div(lax.sub(x, loc), scale) normalize_term = lax.log(lax.div(scale, b)) log_probs = lax.neg( lax.add(normalize_term, lax.mul(lax.add(b, one), lax.log(scaled_x)))) return where(lax.lt(x, lax.add(loc, scale)), -inf, log_probs)
def round(a, decimals=0): if onp.issubdtype(_dtype(a), onp.integer): return a # no-op on integer types if decimals == 0: return lax.round(a) factor = _constant_like(a, 10**decimals) return lax.div(lax.round(lax.mul(a, factor)), factor)
def threefry_random_bits(key: jnp.ndarray, bit_width, shape): """Sample uniform random bits of given width and shape using PRNG key.""" if not _is_threefry_prng_key(key): raise TypeError("_random_bits got invalid prng key.") if bit_width not in (8, 16, 32, 64): raise TypeError("requires 8-, 16-, 32- or 64-bit field width.") shape = core.as_named_shape(shape) for name, size in shape.named_items: real_size = lax.psum(1, name) if real_size != size: raise ValueError( f"The shape of axis {name} was specified as {size}, " f"but it really is {real_size}") axis_index = lax.axis_index(name) key = threefry_fold_in(key, axis_index) size = prod(shape.positional) # Compute ceil(bit_width * size / 32) in a way that is friendly to shape # polymorphism max_count, r = divmod(bit_width * size, 32) if r > 0: max_count += 1 if core.is_constant_dim(max_count): nblocks, rem = divmod(max_count, jnp.iinfo(np.uint32).max) else: nblocks, rem = 0, max_count if not nblocks: bits = threefry_2x32(key, lax.iota(np.uint32, rem)) else: keys = threefry_split(key, nblocks + 1) subkeys, last_key = keys[:-1], keys[-1] blocks = vmap(threefry_2x32, in_axes=(0, None))(subkeys, lax.iota(np.uint32, jnp.iinfo(np.uint32).max)) last = threefry_2x32(last_key, lax.iota(np.uint32, rem)) bits = lax.concatenate([blocks.ravel(), last], 0) dtype = UINT_DTYPES[bit_width] if bit_width == 64: bits = [lax.convert_element_type(x, dtype) for x in jnp.split(bits, 2)] bits = lax.shift_left(bits[0], dtype(32)) | bits[1] elif bit_width in [8, 16]: # this is essentially bits.view(dtype)[:size] bits = lax.bitwise_and( np.uint32(np.iinfo(dtype).max), lax.shift_right_logical( lax.broadcast(bits, (1, )), lax.mul( np.uint32(bit_width), lax.broadcasted_iota(np.uint32, (32 // bit_width, 1), 0)))) bits = lax.reshape(bits, (np.uint32(max_count * 32 // bit_width), ), (1, 0)) bits = lax.convert_element_type(bits, dtype)[:size] return lax.reshape(bits, shape)
def _power(x1, x2): x1, x2 = _promote_args("power", x1, x2) dtype = dtypes.dtype(x1) if not dtypes.issubdtype(dtype, np.integer): return lax.pow(x1, x2) # Integer power => use binary exponentiation. # TODO(phawkins): add integer pow support to XLA. bits = 6 # Anything more would overflow for any x1 > 1 zero = _constant_like(x2, 0) one = _constant_like(x2, 1) # Initialize acc carefully such that pow(0, x2) is zero for x2 != 0 acc = _where(lax.bitwise_and(lax.eq(x1, zero), lax.ne(x2, zero)), zero, one) for _ in range(bits): acc = _where(lax.bitwise_and(x2, one), lax.mul(acc, x1), acc) x1 = lax.mul(x1, x1) x2 = lax.shift_right_logical(x2, one) return acc
def cfun(x): def inner_cond(x): return lax.cond( lax.lt(x, 5), x, lambda x: lax.mul(3, x), 4, lambda y: lax.mul(y, x), ) return lax.cond(lax.lt(x, 2), x, lambda x: lax.mul(2, x), x, inner_cond)