def project(self, x: ep.Tensor, x0: ep.Tensor,
epsilon: float) -> ep.Tensor:
# based on https://github.com/ftramer/MultiRobustness/blob/ad41b63235d13b1b2a177c5f270ab9afa74eee69/pgd_attack.py#L110
delta = flatten(x - x0)
norms = delta.norms.l1(axis=-1)
if (norms <= epsilon).all():
return x
n, d = delta.shape
abs_delta = abs(delta)
mu = -ep.sort(-abs_delta, axis=-1)
cumsums = mu.cumsum(axis=-1)
js = 1.0 / ep.arange(x, 1, d + 1).astype(x.dtype)
temp = mu - js * (cumsums - epsilon)
guarantee_first = ep.arange(x, d).astype(x.dtype) / d
# guarantee_first are small values (< 1) that we add to the boolean
# tensor (only 0 and 1) to break the ties and always return the first
# argmin, i.e. the first value where the boolean tensor is 0
# (otherwise, this is not guaranteed on GPUs, see e.g. PyTorch)
rho = ep.argmin((temp > 0).astype(x.dtype) + guarantee_first, axis=-1)
theta = 1.0 / (1 + rho.astype(x.dtype)) * (cumsums[range(n), rho] -
epsilon)
delta = delta.sign() * ep.maximum(abs_delta - theta[..., ep.newaxis],
0)
delta = delta.reshape(x.shape)
return x0 + delta
def project_onto_l1_ball(x: ep.Tensor, eps: ep.Tensor) -> ep.Tensor:
"""Computes Euclidean projection onto the L1 ball for a batch. [#Duchi08]_
Adapted from the pytorch version by Tony Duan:
https://gist.github.com/tonyduan/1329998205d88c566588e57e3e2c0c55
Args:
x: Batch of arbitrary-size tensors to project, possibly on GPU
eps: radius of l-1 ball to project onto
References:
..[#Duchi08] Efficient Projections onto the l1-Ball for Learning in High Dimensions
John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra.
International Conference on Machine Learning (ICML 2008)
"""
original_shape = x.shape
x = flatten(x)
mask = (ep.norms.l1(x, axis=1) <= eps).astype(x.dtype).expand_dims(1)
mu = ep.flip(ep.sort(ep.abs(x)), axis=-1).astype(x.dtype)
cumsum = ep.cumsum(mu, axis=-1)
arange = ep.arange(x, 1, x.shape[1] + 1).astype(x.dtype)
rho = (ep.max(
((mu * arange >
(cumsum - eps.expand_dims(1)))).astype(x.dtype) * arange,
axis=-1,
) - 1)
# samples already under norm will have to select
rho = ep.maximum(rho, 0)
theta = (cumsum[ep.arange(x, x.shape[0]),
rho.astype(ep.arange(x, 1).dtype)] - eps) / (rho + 1.0)
proj = (ep.abs(x) - theta.expand_dims(1)).clip(min_=0, max_=ep.inf)
x = mask * x + (1 - mask) * proj * ep.sign(x)
return x.reshape(original_shape)
def l2_clipping_aware_rescaling(x,
delta,
eps: float,
a: float = 0.0,
b: float = 1.0): # type: ignore
"""Calculates eta such that norm(clip(x + eta * delta, a, b) - x) == eps.
Assumes x and delta have a batch dimension and eps, a, b, and p are
scalars. If the equation cannot be solved because eps is too large, the
left hand side is maximized.
Args:
x: A batch of inputs (PyTorch Tensor, TensorFlow Eager Tensor, NumPy
Array, JAX Array, or EagerPy Tensor).
delta: A batch of perturbation directions (same shape and type as x).
eps: The target norm (non-negative float).
a: The lower bound of the data domain (float).
b: The upper bound of the data domain (float).
Returns:
eta: A batch of scales with the same number of dimensions as x but all
axis == 1 except for the batch dimension.
"""
(x, delta), restore_fn = ep.astensors_(x, delta)
N = x.shape[0]
assert delta.shape[0] == N
rows = ep.arange(x, N)
delta2 = delta.square().reshape((N, -1))
space = ep.where(delta >= 0, b - x, x - a).reshape((N, -1))
f2 = space.square() / ep.maximum(delta2, 1e-20)
ks = ep.argsort(f2, axis=-1)
f2_sorted = f2[rows[:, ep.newaxis], ks]
m = ep.cumsum(delta2[rows[:, ep.newaxis],
ks.flip(axis=1)], axis=-1).flip(axis=1)
dx = f2_sorted[:, 1:] - f2_sorted[:, :-1]
dx = ep.concatenate((f2_sorted[:, :1], dx), axis=-1)
dy = m * dx
y = ep.cumsum(dy, axis=-1)
c = y >= eps**2
# work-around to get first nonzero element in each row
f = ep.arange(x, c.shape[-1], 0, -1)
j = ep.argmax(c.astype(f.dtype) * f, axis=-1)
eta2 = f2_sorted[rows, j] - (y[rows, j] - eps**2) / m[rows, j]
# it can happen that for certain rows even the largest j is not large enough
# (i.e. c[:, -1] is False), then we will just use it (without any correction) as it's
# the best we can do (this should also be the only cases where m[j] can be
# 0 and they are thus not a problem)
eta2 = ep.where(c[:, -1], eta2, f2_sorted[:, -1])
eta = ep.sqrt(eta2)
eta = eta.reshape((-1, ) + (1, ) * (x.ndim - 1))
# xp = ep.clip(x + eta * delta, a, b)
# l2 = (xp - x).reshape((N, -1)).square().sum(axis=-1).sqrt()
return restore_fn(eta)
def make_uv(t):
'''
Generates UV coordinates
Returns tensors of x coords and y coords each with shape matching t
'''
uvx = ep.expand_dims(ep.arange(t, 0.0, t.shape[1], 1),
axis=0).tile([t.shape[0], 1])
uvy = ep.expand_dims(ep.arange(t, 0.0, t.shape[0], 1),
axis=0).tile([t.shape[1], 1]).transpose()
return uvx, uvy
def test_crossentropy_raises(dummy: Tensor) -> None:
t = ep.arange(dummy, 50).reshape((10, 5)).float32()
t = t / t.max()
ep.crossentropy(t, t.argmax(axis=-1))
t = ep.arange(dummy, 150).reshape((10, 5, 3)).float32()
t = t / t.max()
with pytest.raises(ValueError):
ep.crossentropy(t, t.argmax(axis=-1))
t = ep.arange(dummy, 50).reshape((10, 5)).float32()
t = t / t.max()
with pytest.raises(ValueError):
ep.crossentropy(t, t.argmax(axis=-1)[:8])
def test_matmul_raise(dummy: Tensor) -> None:
t = ep.arange(dummy, 8).float32().reshape((2, 4))
ep.matmul(t, t.T)
with pytest.raises(ValueError):
ep.matmul(t, t[0])
with pytest.raises(ValueError):
ep.matmul(t[0], t)
with pytest.raises(ValueError):
ep.matmul(t[0], t[0])
def project_onto_l1_ball(x: ep.Tensor, eps: ep.Tensor):
"""
Compute Euclidean projection onto the L1 ball for a batch.
min ||x - u||_2 s.t. ||u||_1 <= eps
Inspired by the corresponding numpy version by Adrien Gaidon.
Adapted from the pytorch version by Tony Duan: https://gist.github.com/tonyduan/1329998205d88c566588e57e3e2c0c55
Parameters
----------
x: (batch_size, *) torch array
batch of arbitrary-size tensors to project, possibly on GPU
eps: float
radius of l-1 ball to project onto
Returns
-------
u: (batch_size, *) torch array
batch of projected tensors, reshaped to match the original
Notes
-----
The complexity of this algorithm is in O(dlogd) as it involves sorting x.
References
----------
[1] Efficient Projections onto the l1-Ball for Learning in High Dimensions
John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra.
International Conference on Machine Learning (ICML 2008)
"""
original_shape = x.shape
x = flatten(x)
mask = (ep.norms.l1(x, axis=1) < eps).astype(x.dtype).expand_dims(1)
mu = ep.flip(ep.sort(ep.abs(x)), axis=-1)
cumsum = ep.cumsum(mu, axis=-1)
arange = ep.arange(x, 1, x.shape[1] + 1)
rho = ep.max(
(mu * arange > (cumsum - eps.expand_dims(1))) * arange, axis=-1) - 1
theta = (cumsum[ep.arange(x, x.shape[0]), rho] - eps) / (rho + 1.0)
proj = (ep.abs(x) - theta.expand_dims(1)).clip(min_=0, max_=ep.inf)
x = mask * x + (1 - mask) * proj * ep.sign(x)
return x.reshape(original_shape)
def test_value_and_grad_multiple_args(dummy: Tensor) -> None:
if isinstance(dummy, ep.NumPyTensor):
pytest.skip()
def f(x: Tensor, y: Tensor) -> Tensor:
return (x * y).sum()
t = ep.arange(dummy, 8).float32().reshape((2, 4))
v, g = ep.value_and_grad(f, t, t)
assert v.item() == 140
assert (g == t).all()
def test_pad_raises(dummy: Tensor) -> None:
t = ep.arange(dummy, 120).reshape((2, 3, 4, 5)).float32()
ep.pad(t, ((0, 0), (0, 0), (2, 3), (1, 2)), mode="constant")
with pytest.raises(ValueError):
ep.pad(t, ((0, 0), (2, 3), (1, 2)), mode="constant")
with pytest.raises(ValueError):
ep.pad(
t, ((0, 0), (0, 0, 1, 2), (2, 3), (1, 2)), mode="constant" # type: ignore
)
with pytest.raises(ValueError):
ep.pad(t, ((0, 0), (0, 0), (2, 3), (1, 2)), mode="foo")
def test_value_and_grad_fn(dummy: Tensor) -> None:
if isinstance(dummy, ep.NumPyTensor):
pytest.skip()
def f(x: ep.Tensor) -> ep.Tensor:
return x.square().sum()
vgf = ep.value_and_grad_fn(dummy, f)
t = ep.arange(dummy, 8).float32().reshape((2, 4))
v, g = vgf(t)
assert v.item() == 140
assert (g == 2 * t).all()
def project(self, x: ep.Tensor, x0: ep.Tensor,
epsilon: ep.Tensor) -> ep.Tensor:
flatten_delta = flatten(x - x0)
n, d = flatten_delta.shape
abs_delta = abs(flatten_delta)
epsilon = epsilon.astype(ep.arange(x, 1).dtype)
rows = range(n)
idx_sorted = ep.flip(ep.argsort(abs_delta, axis=1), -1)[rows, epsilon]
thresholds = (ep.ones_like(flatten_delta).T *
abs_delta[rows, idx_sorted]).T
clipped = ep.where(abs_delta >= thresholds, flatten_delta, 0)
return x0 + clipped.reshape(x0.shape).astype(x0.dtype)
def test_value_aux_and_grad(dummy: Tensor) -> None:
if isinstance(dummy, ep.NumPyTensor):
pytest.skip()
def f(x: Tensor) -> Tuple[Tensor, Tensor]:
x = x.square()
return x.sum(), x
t = ep.arange(dummy, 8).float32().reshape((2, 4))
v, aux, g = ep.value_aux_and_grad(f, t)
assert v.item() == 140
assert (aux == t.square()).all()
assert (g == 2 * t).all()
def test_tensorboard(logdir: Union[Literal[False], None, str], tmp_path: Any,
dummy: ep.Tensor) -> None:
if logdir == "temp":
logdir = tmp_path
if logdir:
before = len(list(tmp_path.iterdir()))
tb = fbn.tensorboard.TensorBoard(logdir)
tb.scalar("a_scalar", 5, step=1)
x = ep.ones(dummy, 10)
tb.mean("a_mean", x, step=2)
x = ep.ones(dummy, 10) == ep.arange(dummy, 10)
tb.probability("a_probability", x, step=2)
x = ep.arange(dummy, 10).float32()
cond = ep.ones(dummy, 10) == (ep.arange(dummy, 10) % 2)
tb.conditional_mean("a_conditional_mean", x, cond, step=2)
x = ep.arange(dummy, 10).float32()
cond = ep.ones(dummy, 10) == ep.zeros(dummy, 10)
tb.conditional_mean("a_conditional_mean_false", x, cond, step=2)
x = ep.ones(dummy, 10) == ep.arange(dummy, 10)
y = ep.ones(dummy, 10) == (ep.arange(dummy, 10) % 2)
tb.probability_ratio("a_probability_ratio", x, y, step=5)
x = ep.ones(dummy, 10) == (ep.arange(dummy, 10) % 2)
y = ep.ones(dummy, 10) == ep.zeros(dummy, 10)
tb.probability_ratio("a_probability_ratio_y_zero", x, y, step=5)
x = ep.arange(dummy, 10).float32()
tb.histogram("a_histogram", x, step=9, first=False)
tb.histogram("a_histogram", x, step=10, first=True)
tb.close()
if logdir:
after = len(list(tmp_path.iterdir()))
assert after > before # make sure something has been written
def test_onehot_like_raises(dummy: Tensor) -> None:
t = ep.arange(dummy, 18).float32().reshape((6, 3))
indices = ep.arange(t, 6) // 2
ep.onehot_like(t, indices)
t = ep.arange(dummy, 90).float32().reshape((6, 3, 5))
indices = ep.arange(t, 6) // 2
with pytest.raises(ValueError):
ep.onehot_like(t, indices)
t = ep.arange(dummy, 18).float32().reshape((6, 3))
indices = ep.arange(t, 6).reshape((6, 1)) // 2
with pytest.raises(ValueError):
ep.onehot_like(t, indices)
t = ep.arange(dummy, 18).float32().reshape((6, 3))
indices = ep.arange(t, 5) // 2
with pytest.raises(ValueError):
ep.onehot_like(t, indices)
def test_getitem_slice_slice(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 32).float32().reshape((4, 8))
return t[:, :3]
def test_getitem_tuple_list_tensor(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 32).float32().reshape((8, 4))
rows = list(range(len(t)))
indices = ep.arange(t, len(t)) % t.shape[1]
return t[rows, indices]
def test_getitem_tuple_range_range(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 36).float32().reshape((6, 6))
rows = cols = range(len(t))
return t[rows, cols]
def test_getitem_ndarray(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 32).float32()
indices = np.arange(3, 10, 2)
return t[indices]
def test_getitem_list(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 32).float32()
indices = list(range(3, 10, 2))
return t[indices]
def test_getitem_tensor(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 32).float32()
indices = ep.arange(t, 3, 10, 2)
return t[indices]
def test_getitem_ellipsis_newaxis(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 8).float32()
return t[..., ep.newaxis]
def test_getitem_boolean_tensor(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 32).float32().reshape((4, 8))
indices = ep.arange(t, 4) <= 2
return t[indices]
def test_prod_int(t: Tensor) -> Tensor:
return ep.prod(ep.arange(t, 5))
def test_take_along_axis_3d(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 64).float32().reshape((2, 8, 4))
indices = ep.arange(t, 2 * 8).reshape((2, 8, 1)) % t.shape[-1]
return ep.take_along_axis(t, indices, axis=-1)
def test_take_along_axis_2d(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 32).float32().reshape((8, 4))
indices = ep.arange(t, len(t)) % t.shape[-1]
return ep.take_along_axis(t, indices[..., ep.newaxis], axis=-1)
def test_getitem_tuple(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 8).float32().reshape((2, 4))
return t[1, 3]
def test_sum_int(t: Tensor) -> Tensor:
return ep.sum(ep.arange(t, 5))
def test_take_along_axis_2d_first_raises(dummy: Tensor) -> None:
t = ep.arange(dummy, 32).float32().reshape((8, 4))
indices = ep.arange(t, t.shape[-1]) % t.shape[0]
with pytest.raises(NotImplementedError):
ep.take_along_axis(t, indices[ep.newaxis], axis=0)
def test_any_axes(dummy: Tensor) -> Tensor:
t = ep.arange(dummy, 30).float32().reshape((3, 5, 2))
return ep.any(t > 3, axis=(0, 1))
def test_format(dummy: Tensor) -> bool:
t = ep.arange(dummy, 5).sum()
return f"{t:.1f}" == "10.0"