def gibbs_loop(i, rng_particles_log_posteriors):
rng, particles, log_posteriors = rng_particles_log_posteriors
i = i % num_patients
# flip values at index i
particles_flipped = jax.ops.index_update(
particles, jax.ops.index[:, i], np.logical_not(particles[:, i]))
# compute log_posterior of flipped particles
log_posteriors_flipped_at_i = bayes.tempered_logpos_logbase(
particles_flipped, log_posterior_params, log_base_measure_params,
rho)
# compute acceptance probability, depending on whether we use Liu mod.
if liu_modification:
log_proposal_ratio = log_posteriors_flipped_at_i - log_posteriors
else:
log_proposal_ratio = log_posteriors_flipped_at_i - np.logaddexp(
log_posteriors_flipped_at_i, log_posteriors)
# here the MH thresholding is implicitly done.
rng, rng_unif = jax.random.split(rng, 2)
random_values = jax.random.uniform(rng_unif, particles.shape[:1])
flipped_at_i = np.log(random_values) < log_proposal_ratio
selected_at_i = np.logical_xor(flipped_at_i, particles[:, i])
particles = jax.ops.index_update(particles, jax.ops.index[:, i],
selected_at_i)
log_posteriors = np.where(flipped_at_i, log_posteriors_flipped_at_i,
log_posteriors)
return [rng, particles, log_posteriors]
def matte(vis, acc, dark=0.8, light=1.0, width=8):
"""Set non-accumulated pixels to a Photoshop-esque checker pattern."""
bg_mask = jnp.logical_xor(
(jnp.arange(acc.shape[0]) % (2 * width) // width)[:, None],
(jnp.arange(acc.shape[1]) % (2 * width) // width)[None, :])
bg = jnp.where(bg_mask, light, dark)
return vis * acc[:, :, None] + (bg * (1 - acc))[:, :, None]
def __call__(self, rng, state):
"""Produces new groups from 1st wave of results & adds them to stack."""
num_rows = 8
num_cols = 12
# check this is the first time selector is called (or not used)
if state.past_groups.shape[0] == num_rows + num_cols:
# sum groups that have returned positive
returned_positive = np.sum(
state.past_groups[state.past_test_results, :], axis=0)
total_positive_first_block = np.sum(state.past_test_results[0:num_rows],
axis=0)
total_positive_second_block = np.sum(
state.past_test_results[num_rows:num_rows+num_cols], axis=0)
# check if incoherence in row/columns
if np.logical_xor(total_positive_first_block > 0,
total_positive_second_block > 0):
# test everyone that was ever included in a positive group
reflex_to_test, = np.where(returned_positive)
else:
# test individually those that returned positive at least twice.
reflex_to_test, = np.where(returned_positive > 1)
new_groups = jax.nn.one_hot(reflex_to_test, state.num_patients).astype(bool)
state.add_groups_to_test(new_groups)
logging.warning('Added %i groups to test', new_groups.shape[0])
logging.debug(new_groups.astype(np.int32))
else:
state.all_cleared = True
return state
def body(args, _):
"""Body for the while loop executing the binary search."""
bit_index, value = args
new_value = jnp.bitwise_or(value, jnp.left_shift(1, bit_index))
larger = larger_count(scores, bitcast(new_value, jnp.float32))
next_value = jnp.where(jnp.logical_xor(larger >= k, kth_negative),
new_value, value)
return (bit_index - 1, next_value), None
def compute_opaqueness_mask(weights, depth_threshold=0.5):
"""Computes a mask which will be 1.0 at the depth point.
Args:
weights: the density weights from NeRF.
depth_threshold: the accumulation threshold which will be used as the depth
termination point.
Returns:
A tensor containing a mask with the same size as weights that has one
element long the sample dimension that is 1.0. This element is the point
where the 'surface' is.
"""
cumulative_contribution = jnp.cumsum(weights, axis=-1)
depth_threshold = jnp.array(depth_threshold, dtype=weights.dtype)
opaqueness = cumulative_contribution >= depth_threshold
false_padding = jnp.zeros_like(opaqueness[..., :1])
padded_opaqueness = jnp.concatenate([false_padding, opaqueness[..., :-1]], axis=-1)
opaqueness_mask = jnp.logical_xor(opaqueness, padded_opaqueness)
opaqueness_mask = opaqueness_mask.astype(weights.dtype)
return opaqueness_mask
def SmoothedBoxPrior(theta_dim=5,
lower=0.0,
upper=1.0,
sigma=0.1,
variance=False):
assert np.all(lower < upper), "lower must be less than upper"
assert np.all(sigma > 0), "sigma must be greater than zero"
assert np.logical_xor(sigma,
variance), "specify only one of sigma and variance"
if not variance:
variance = sigma**2
_center = (upper + lower) / 2.0
_range = (upper - lower) / 2.0
def log_prob(theta):
"""Inspired by SmoothedBoxPrior From GPyTorch
If theta is inside the bounds, return constant.
If theta is outside the bounds, return log prob from sharp normal
Can accomplish this saying the distance from the edges of the theta range
is sampled from a normal distribution (clipped at zero to not go negative)
"""
_theta_dist = np.clip(np.abs(theta - _center) - _range, 0, None)
return -0.5 * (_theta_dist**2 / variance +
np.log(2 * np.pi * variance))
def sample(rng, num_samples: int = 1):
"""
Samples are taken from a hard uniform distribution between the bounds
"""
return jax.random.uniform(rng,
shape=(num_samples, theta_dim),
minval=lower,
maxval=upper)
return log_prob, sample
def test_xor(self):
self.binary_check(lambda x, y: np.logical_xor(x, y))
def logical_xor(x1, x2): if isinstance(x1, JaxArray): x1 = x1.value if isinstance(x2, JaxArray): x2 = x2.value return JaxArray(jnp.logical_xor(x1, x2))
def onnx_xor(a, b):
return jnp.logical_xor(a, b)
np.subtract( # pylint: disable=g-long-lambda
logits,
reduce_logsumexp(
logits, -1 if axis is None else axis, keepdims=True))))
logical_and = utils.copy_docstring(
tf.math.logical_and, lambda x, y, name=None: np.logical_and(x, y))
logical_not = utils.copy_docstring(tf.math.logical_not,
lambda x, name=None: np.logical_not(x))
logical_or = utils.copy_docstring(tf.math.logical_or,
lambda x, y, name=None: np.logical_or(x, y))
logical_xor = utils.copy_docstring(
tf.math.logical_xor, lambda x, y, name=None: np.logical_xor(x, y))
maximum = utils.copy_docstring(tf.math.maximum,
lambda x, y, name=None: np.maximum(x, y))
minimum = utils.copy_docstring(tf.math.minimum,
lambda x, y, name=None: np.minimum(x, y))
multiply = utils.copy_docstring(tf.math.multiply,
lambda x, y, name=None: np.multiply(x, y))
multiply_no_nan = utils.copy_docstring(
tf.math.multiply_no_nan,
lambda x, y, name=None: np.where( # pylint: disable=g-long-lambda
onp.broadcast_to(np.equal(y, 0.),
np.array(x).shape), np.zeros_like(np.multiply(x, y)),
def test_xor(self): self.binary_check(lambda x, y: np.logical_xor(x, y))
@jtu.skip_on_devices("tpu")
def test_xor(self): self.binary_check(lambda x, y: np.logical_xor(x, y)) def test_process_call(self):