File size: 4,140 Bytes
62d3300 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
"""Utils for geometry library."""
from collections.abc import Iterable
import numbers
import jax
from jax import lax
import jax.numpy as jnp
def safe_select(condition, true_fn, false_fn):
"""Safe version of selection (i.e. `where`).
This applies the double-where trick.
Like jnp.where, this function will still execute both branches and is
expected to be more lightweight than lax.cond. Other than NaN-semantics,
safe_select(condition, true_fn, false_fn) is equivalent to
jax.tree.map(lambda x, y: jnp.where(condition, x, y),
true_fn(),
false_fn()),
Compared to the naive implementation above, safe_select provides the
following guarantee: in either the forward or backward pass, a NaN produced
*during the execution of true_fn()* will not propagate to the rest of the
computation and similarly for false_fn. It is very important to note that
while true_fn and false_fn will typically close over other tensors (i.e. they
use values computed prior to the safe_select function), there is no NaN-safety
for the backward pass of closed over values. It is important than any NaN's
are produced within the branch functions and not before them. For example,
safe_select(x < eps, lambda: 0., lambda: jnp.sqrt(x))
will not produce NaN on the backward pass even if x == 0. since sqrt happens
within the false_fn, but the very similar
y = jnp.sqrt(x)
safe_select(x < eps, lambda: 0., lambda: y)
will produce a NaN on the backward pass if x == 0 because the sqrt happens
prior to the false_fn.
Args:
condition: Boolean array to use in where
true_fn: Zero-argument function to construct the values used in the True
condition. Tensors that this function closes over will be extracted
automatically to implement the double-where trick to suppress spurious NaN
propagation.
false_fn: False branch equivalent of true_fn
Returns:
Resulting PyTree equivalent to tree_map line above.
"""
true_fn, true_args = jax.closure_convert(true_fn)
false_fn, false_args = jax.closure_convert(false_fn)
true_args = jax.tree.map(
lambda x: jnp.where(condition, x, lax.stop_gradient(x)), true_args
)
false_args = jax.tree.map(
lambda x: jnp.where(condition, lax.stop_gradient(x), x), false_args
)
return jax.tree.map(
lambda x, y: jnp.where(condition, x, y),
true_fn(*true_args),
false_fn(*false_args),
)
def unstack(value: jnp.ndarray, axis: int = -1) -> list[jnp.ndarray]:
return [
jnp.squeeze(v, axis=axis)
for v in jnp.split(value, value.shape[axis], axis=axis)
]
def angdiff(alpha: jnp.ndarray, beta: jnp.ndarray) -> jnp.ndarray:
"""Compute absolute difference between two angles."""
d = alpha - beta
d = (d + jnp.pi) % (2 * jnp.pi) - jnp.pi
return d
def safe_arctan2(
x1: jnp.ndarray, x2: jnp.ndarray, eps: float = 1e-8
) -> jnp.ndarray:
"""Safe version of arctan2 that avoids NaN gradients when x1=x2=0."""
return safe_select(
jnp.abs(x1) + jnp.abs(x2) < eps,
lambda: jnp.zeros_like(jnp.arctan2(x1, x2)),
lambda: jnp.arctan2(x1, x2),
)
def weighted_mean(
*,
weights: jnp.ndarray,
value: jnp.ndarray,
axis: int | Iterable[int] | None = None,
eps: float = 1e-10,
) -> jnp.ndarray:
"""Computes weighted mean in a safe way that avoids NaNs.
This is equivalent to jnp.average for the case eps=0.0, but adds a small
constant to the denominator of the weighted average to avoid NaNs.
'weights' should be broadcastable to the shape of value.
Args:
weights: Weights to weight value by.
value: Values to average
axis: Axes to average over.
eps: Epsilon to add to the denominator.
Returns:
Weighted average.
"""
weights = jnp.asarray(weights, dtype=value.dtype)
weights = jnp.broadcast_to(weights, value.shape)
weights_shape = weights.shape
if isinstance(axis, numbers.Integral):
axis = [axis]
elif axis is None:
axis = list(range(len(weights_shape)))
return jnp.sum(weights * value, axis=axis) / (
jnp.sum(weights, axis=axis) + eps
)
|