def test_seq_mult(shape_1: Tuple[int, ...], num_arrays: int,
data: st.DataObject):
shape_2 = data.draw(hnp.broadcastable_shapes(shape_1), label="shape_2")
shapes = [shape_1, shape_2]
pair = shapes
for i in range(num_arrays):
# ensure sequence of shapes is mutually-broadcastable
broadcasted = _broadcast_shapes(*pair)
shapes.append(
data.draw(hnp.broadcastable_shapes(broadcasted),
label="shape_{}".format(i + 3)))
pair = [broadcasted, shapes[-1]]
tensors = [
Tensor(
data.draw(
hnp.arrays(shape=shape,
dtype=np.float32,
elements=st.floats(-10, 10, width=32))))
for shape in shapes
]
note("tensors: {}".format(tensors))
tensors_copy = [x.copy() for x in tensors]
f = multiply_sequence(*tensors)
f1 = reduce(lambda x, y: x * y,
(var for n, var in enumerate(tensors_copy)))
assert_allclose(f.data, f1.data)
f.sum().backward()
f1.sum().backward()
assert_allclose(f.data, f1.data, rtol=1e-4, atol=1e-4)
for n, (expected, actual) in enumerate(zip(tensors_copy, tensors)):
assert_allclose(
expected.grad,
actual.grad,
rtol=1e-3,
atol=1e-3,
err_msg="tensor-{}".format(n),
)
f.null_gradients()
assert all(x.grad is None for x in tensors)
assert all(not x._ops for x in tensors)
def test_broadcastable_shape_has_good_default_values(shape, data):
# This test ensures that default parameters can always produce broadcast-compatible shapes
broadcastable_shape = data.draw(
nps.broadcastable_shapes(shape), label="broadcastable_shapes"
)
# error if drawn shape for b is not broadcast-compatible
_broadcast_shapes(shape, broadcastable_shape)
def test_broadcastable_shape_bounds_are_satisfied(shape, data):
min_dims = data.draw(st.integers(0, 32), label="min_dims")
max_dims = data.draw(st.none() | st.integers(min_dims, 32), label="max_dims")
min_side = data.draw(st.integers(0, 3), label="min_side")
max_side = data.draw(st.none() | st.integers(min_side, 6), label="max_side")
try:
bshape = data.draw(
nps.broadcastable_shapes(
shape,
min_side=min_side,
max_side=max_side,
min_dims=min_dims,
max_dims=max_dims,
),
label="bshape",
)
except InvalidArgument:
raise UnsatisfiedAssumption from None
if max_dims is None:
max_dims = max(len(shape), min_dims) + 2
if max_side is None:
max_side = max(tuple(shape[::-1][:max_dims]) + (min_side,)) + 2
assert isinstance(bshape, tuple) and all(isinstance(s, int) for s in bshape)
assert min_dims <= len(bshape) <= max_dims
assert all(min_side <= s <= max_side for s in bshape)
def test_broadcastable_shape_bounds_are_satisfied(shape, data):
min_dim = data.draw(st.integers(0, 4), label="min_dim")
max_dim = data.draw(st.one_of(st.none(), st.integers(min_dim, 4)),
label="max_dim")
min_side = data.draw(st.integers(0, 3), label="min_side")
max_side = data.draw(st.one_of(st.none(), st.integers(min_side, 6)),
label="max_side")
try:
bshape = data.draw(
nps.broadcastable_shapes(
shape,
min_side=min_side,
max_side=max_side,
min_dims=min_dim,
max_dims=max_dim,
),
label="bshape",
)
except InvalidArgument:
assume(False)
if max_dim is None:
max_dim = max(len(shape), min_dim) + 2
if max_side is None:
max_side = max(tuple(shape[::-1][:max_dim]) + (min_side, )) + 2
assert isinstance(bshape, tuple) and all(
isinstance(s, int) for s in bshape)
assert min_dim <= len(bshape) <= max_dim
assert all(min_side <= s <= max_side for s in bshape)
def gen_broadcast_to(draw):
shape1 = draw(
st.lists(st.integers(min_value=2, max_value=5), min_size=2, max_size=3).map(
tuple
)
)
shape2 = draw(broadcastable_shapes(shape1, min_dims=2, max_dims=3))
return shape1, shape2
def to_min_max(arr: np.ndarray) -> st.SearchStrategy:
bnd_shape = hnp.broadcastable_shapes(
shape=arr.shape,
max_dims=arr.ndim,
max_side=min(arr.shape) if arr.ndim else 1)
bnd_strat = hnp.arrays(shape=bnd_shape,
elements=st.floats(-1e6, 1e6),
dtype=np.float64)
return st.fixed_dictionaries(dict(a_min=bnd_strat, a_max=bnd_strat))
def test_broadcastable_shape_can_generate_arbitrary_ndims(
shape, max_dims, data):
# ensures that generates shapes can possess any length in [min_dims, max_dims]
desired_ndim = data.draw(st.integers(0, max_dims), label="desired_ndim")
min_dims = data.draw(st.one_of(st.none(), st.integers(0, desired_ndim)),
label="min_dims")
args = (dict(min_dims=min_dims) if min_dims is not None else {}
) # check default arg behavior too
find_any(
nps.broadcastable_shapes(shape, min_side=0, max_dims=max_dims, **args),
lambda x: len(x) == desired_ndim,
)
def test_broadcastable_shape_adjusts_max_dim_with_explicit_bounds(max_dims, data):
# Ensures that `broadcastable_shapes` limits itself to satisfiable dimensions
# Broadcastable values can only be drawn for dims 0-3 for these shapes
shape = data.draw(st.sampled_from([(5, 3, 2, 1), (0, 3, 2, 1)]), label="shape")
broadcastable_shape = data.draw(
nps.broadcastable_shapes(
shape, min_side=2, max_side=3, min_dims=3, max_dims=max_dims
),
label="broadcastable_shapes",
)
assert len(broadcastable_shape) == 3
# error if drawn shape for b is not broadcast-compatible
_broadcast_shapes(shape, broadcastable_shape)
def test_broadcastable_shape_can_broadcast(min_dims, shape, data):
max_dims = data.draw(st.none() | st.integers(min_dims, 32), label="max_dims")
min_side, max_side = _draw_valid_bounds(data, shape, max_dims)
broadcastable_shape = data.draw(
nps.broadcastable_shapes(
shape,
min_side=min_side,
max_side=max_side,
min_dims=min_dims,
max_dims=max_dims,
),
label="broadcastable_shapes",
)
# error if drawn shape for b is not broadcast-compatible
_broadcast_shapes(shape, broadcastable_shape)
def test_broadcastable_shape_shrinking_with_singleton_out_of_bounds(
min_dims, min_side, data
):
max_dims = data.draw(st.none() | st.integers(min_dims, 32), label="max_dims")
max_side = data.draw(st.none() | st.integers(min_side, 6), label="max_side")
shape = data.draw(st.integers(1, 4).map(lambda n: n * (1,)), label="shape")
smallest = minimal(
nps.broadcastable_shapes(
shape,
min_side=min_side,
max_side=max_side,
min_dims=min_dims,
max_dims=max_dims,
)
)
assert smallest == (min_side,) * min_dims
def test_broadcastable_shape_can_broadcast(min_dim, shape, data):
max_dim = data.draw(st.one_of(st.none(), st.integers(min_dim, 5)), label="max_dim")
min_side, max_side = _draw_valid_bounds(data, shape, max_dim)
broadcastable_shape = data.draw(
nps.broadcastable_shapes(
shape,
min_side=min_side,
max_side=max_side,
min_dims=min_dim,
max_dims=max_dim,
),
label="broadcastable_shapes",
)
a = np.zeros(shape, dtype="uint8")
b = np.zeros(broadcastable_shape, dtype="uint8")
np.broadcast(a, b) # error if drawn shape for b is not broadcast-compatible
def test_broadcastable_shape_shrinking_with_singleton_out_of_bounds(
min_dim, min_side, data):
max_dim = data.draw(st.one_of(st.none(), st.integers(min_dim, 4)),
label="max_dim")
max_side = data.draw(st.one_of(st.none(), st.integers(min_side, 6)),
label="max_side")
ndims = data.draw(st.integers(1, 4), label="ndim")
shape = (1, ) * ndims
smallest = minimal(
nps.broadcastable_shapes(
shape,
min_side=min_side,
max_side=max_side,
min_dims=min_dim,
max_dims=max_dim,
))
assert smallest == (min_side, ) * min_dim
def test_minimize_broadcastable_shape(min_dims, shape, data):
# Ensure aligned dimensions of broadcastable shape minimizes to `(1,) * min_dims`
max_dims = data.draw(st.none() | st.integers(min_dims, 32), label="max_dims")
min_side, max_side = _draw_valid_bounds(data, shape, max_dims, permit_none=False)
smallest = minimal(
nps.broadcastable_shapes(
shape,
min_side=min_side,
max_side=max_side,
min_dims=min_dims,
max_dims=max_dims,
)
)
note("(smallest): {}".format(smallest))
n_leading = max(len(smallest) - len(shape), 0)
n_aligned = max(len(smallest) - n_leading, 0)
expected = [min_side] * n_leading + [
1 if min_side <= 1 <= max_side else i for i in shape[len(shape) - n_aligned :]
]
assert tuple(expected) == smallest
assert len(smallest) == min_size and all(k > -1 for k in smallest)
@settings(deadline=None, max_examples=10)
@given(ndim=st.integers(0, 3), data=st.data())
def test_minimize_negative_tuple_axes(ndim, data):
min_size = data.draw(st.integers(0, ndim), label="min_size")
max_size = data.draw(st.integers(min_size, ndim), label="max_size")
smallest = minimal(
nps.valid_tuple_axes(ndim, min_size=min_size, max_size=max_size),
lambda x: all(i < 0 for i in x),
)
assert len(smallest) == min_size
@given(nps.broadcastable_shapes((), min_side=0, max_side=0, min_dims=0, max_dims=0))
def test_broadcastable_empty_shape(shape):
assert shape == ()
@settings(deadline=None, suppress_health_check=[HealthCheck.too_slow])
@given(shape=ANY_SHAPE, data=st.data())
def test_broadcastable_shape_bounds_are_satisfied(shape, data):
min_dims = data.draw(st.integers(0, 32), label="min_dims")
max_dims = data.draw(st.none() | st.integers(min_dims, 32), label="max_dims")
min_side = data.draw(st.integers(0, 3), label="min_side")
max_side = data.draw(st.none() | st.integers(min_side, 6), label="max_side")
try:
bshape = data.draw(
nps.broadcastable_shapes(
shape,
def condition_strat(*arrs):
shape = np.broadcast(*arrs).shape
return hnp.arrays(shape=hnp.broadcastable_shapes(shape=shape), dtype=bool)
