Repository navigation
Ufuncs and matmul with an out= larger than a symmetric-tagged operand raise ValueError or SymmetryError on calls numpy accepts #296
Description
Summary
When an operand carries a symmetry tag and out= has more dimensions than the operands, numpy broadcasts the result into out, but flopscope checks that result against the operand's group in the operand's own axis numbering. For S, an (n, n) matrix tagged on axes (0, 1), and an out of shape (3, n, n), the check transposes axes 0 and 1 of the (3, n, n) result and compares it with an (n, 3, n) array:
- an untagged
out=of shape(3, n, n)raisesValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16), forexp,add(S, S),multiply(S, 2.0)andmatmul(S, S); - an
out=of shape(1, n, n)or(n, n, n)raisesSymmetryError: ... max deviation = inf: those shapes do broadcast, so the wrong axes are compared; - a tagged
out=whose tag is the symmetry the result really has in that frame, axes(1, 2), raisesValueError: out symmetry does not match result symmetry.
numpy accepts all of these calls (checked on numpy 2.2.6 and 2.3.3). The same values untagged work, an out= of the operand's own shape works, and flopscope itself derives the tag (1, 2) for this frame in exp(broadcast_to(S, (3, n, n))). An fnp.zeros((3, n, n)) destination happens to work, only because its auto-inferred tag makes _prepare_symmetric_out drop the symmetry. The failures are the same in-process and on the client/server stack, on origin/main, and with #288 (the fix for #275) applied.
Repro (flopscope 0.12.0 in-process with numpy 2.3.3, and flopscope-client 0.12.1 against flopscope-server 0.12.1 with numpy 2.2.6; origin/main at 80e72a9 prints the same)
import warnings; warnings.filterwarnings("ignore")
import flopscope as flops
import flopscope.numpy as fnp
def measure(ctx, fn):
"""FLOPs billed for fn() alone, plus its result."""
before = int(ctx.flops_used)
out = fn()
return int(ctx.flops_used) - before, out
def same(a, b):
"""True when both arrays hold identical values (checked outside every measured window)."""
return bool(fnp.array_equal(a, b))
with flops.BudgetContext(flop_budget=2**40, quiet=True) as c:
n = 16
a = fnp.random.default_rng(0).standard_normal((n, n))
P = a + fnp.transpose(a) # symmetric values, untagged
S = flops.as_symmetric(P, symmetry=(0, 1)) # the same values, tagged
cases = (
# (label, call, the same values computed from the untagged P)
("exp(P, out=empty((3, n, n))) [untagged input]", lambda: fnp.exp(P, out=fnp.empty((3, n, n))),
lambda: fnp.exp(fnp.broadcast_to(P, (3, n, n)))),
("exp(S, out=empty((3, n, n)))", lambda: fnp.exp(S, out=fnp.empty((3, n, n))),
lambda: fnp.exp(fnp.broadcast_to(P, (3, n, n)))),
("exp(S, out=standard_normal((3, n, n)))", lambda: fnp.exp(S, out=fnp.random.standard_normal((3, n, n))),
lambda: fnp.exp(fnp.broadcast_to(P, (3, n, n)))),
("exp(S, out=empty((1, n, n)))", lambda: fnp.exp(S, out=fnp.empty((1, n, n))),
lambda: fnp.exp(fnp.broadcast_to(P, (1, n, n)))),
("exp(S, out=empty((n, n, n)))", lambda: fnp.exp(S, out=fnp.empty((n, n, n))),
lambda: fnp.exp(fnp.broadcast_to(P, (n, n, n)))),
("add(S, S, out=empty((3, n, n)))", lambda: fnp.add(S, S, out=fnp.empty((3, n, n))),
lambda: fnp.add(fnp.broadcast_to(P, (3, n, n)), P)),
("multiply(S, 2.0, out=empty((3, n, n)))", lambda: fnp.multiply(S, 2.0, out=fnp.empty((3, n, n))),
lambda: fnp.multiply(fnp.broadcast_to(P, (3, n, n)), 2.0)),
("matmul(S, S, out=empty((3, n, n)))", lambda: fnp.matmul(S, S, out=fnp.empty((3, n, n))),
lambda: fnp.matmul(fnp.broadcast_to(P, (3, n, n)), P)),
("exp(S, out=as_symmetric(zeros((3, n, n)), symmetry=(1, 2)))",
lambda: fnp.exp(S, out=flops.as_symmetric(fnp.zeros((3, n, n)), symmetry=(1, 2))),
lambda: fnp.exp(fnp.broadcast_to(P, (3, n, n)))),
)
for label, call, ref in cases:
try:
b, r = measure(c, call)
except Exception as e:
print(f"{label:62} raises {type(e).__name__}: {e}")
continue
print(f"{label:62} billed {b:>6,} | identical to the untagged spelling {same(r, ref())}")
# controls: the same tagged S with an out= of its own shape, and the tag flopscope derives itself in the wider frame
O1, O2 = fnp.empty((n, n)), flops.as_symmetric(fnp.zeros((n, n)), symmetry=(0, 1))
b1, r1 = measure(c, lambda: fnp.exp(S, out=O1))
b2, r2 = measure(c, lambda: fnp.exp(S, out=O2))
print(f"control: exp(S, out=empty((n, n))) billed {b1:,}; exp(S, out=as_symmetric(zeros((n, n)), symmetry=(0, 1))) "
f"billed {b2:,}, result tag {r2.symmetry}")
b3, r3 = measure(c, lambda: fnp.exp(fnp.broadcast_to(S, (3, n, n))))
print(f"control: exp(broadcast_to(S, (3, n, n))) billed {b3:,}, result tag {r3.symmetry}, "
f"identical {same(r3, fnp.exp(fnp.broadcast_to(P, (3, n, n))))}")
Z = fnp.zeros((3, n, n)) # auto-tagged (1, 2) as an inferred tag
ztag = Z.symmetry
b4, r4 = measure(c, lambda: fnp.exp(S, out=Z))
print(f"control: exp(S, out=zeros((3, n, n))) billed {b4:,} (the zeros carried the inferred tag {ztag}), "
f"identical {same(r4, fnp.exp(fnp.broadcast_to(P, (3, n, n))))}")Output, flopscope 0.12.0 in-process:
exp(P, out=empty((3, n, n))) [untagged input] billed 8,192 | identical to the untagged spelling True
exp(S, out=empty((3, n, n))) raises ValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16)
exp(S, out=standard_normal((3, n, n))) raises ValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16)
exp(S, out=empty((1, n, n))) raises SymmetryError: Tensor not symmetric along axes (0, 1): max deviation = inf (tolerance: atol=1e-06, rtol=1e-05). See: https://aicrowd.github.io/flopscope/docs/guides/symmetry
exp(S, out=empty((n, n, n))) raises SymmetryError: Tensor not symmetric along axes (0, 1): max deviation = inf (tolerance: atol=1e-06, rtol=1e-05). See: https://aicrowd.github.io/flopscope/docs/guides/symmetry
add(S, S, out=empty((3, n, n))) raises ValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16)
multiply(S, 2.0, out=empty((3, n, n))) raises ValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16)
matmul(S, S, out=empty((3, n, n))) raises ValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16)
exp(S, out=as_symmetric(zeros((3, n, n)), symmetry=(1, 2))) raises ValueError: out symmetry does not match result symmetry
control: exp(S, out=empty((n, n))) billed 4,352; exp(S, out=as_symmetric(zeros((n, n)), symmetry=(0, 1))) billed 4,352, result tag SymmetryGroup([1, 0], axes=(0, 1))
control: exp(broadcast_to(S, (3, n, n))) billed 13,056, result tag SymmetryGroup([1, 0], axes=(1, 2)), identical True
control: exp(S, out=zeros((3, n, n))) billed 8,192 (the zeros carried the inferred tag SymmetryGroup([1, 0], axes=(1, 2))), identical True
Output, flopscope-client 0.12.1 against flopscope-server 0.12.1:
exp(P, out=empty((3, n, n))) [untagged input] billed 8,192 | identical to the untagged spelling True
exp(S, out=empty((3, n, n))) raises ValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16)
exp(S, out=standard_normal((3, n, n))) raises ValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16)
exp(S, out=empty((1, n, n))) raises SymmetryError: Tensor not symmetric along axes (0, 1): max deviation = inf (tolerance: atol=1e-06, rtol=1e-05). See: https://aicrowd.github.io/flopscope/docs/guides/symmetry
exp(S, out=empty((n, n, n))) raises SymmetryError: Tensor not symmetric along axes (0, 1): max deviation = inf (tolerance: atol=1e-06, rtol=1e-05). See: https://aicrowd.github.io/flopscope/docs/guides/symmetry
add(S, S, out=empty((3, n, n))) raises ValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16)
multiply(S, 2.0, out=empty((3, n, n))) raises ValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16)
matmul(S, S, out=empty((3, n, n))) raises ValueError: operands could not be broadcast together with shapes (3,16,16) (16,3,16)
exp(S, out=as_symmetric(zeros((3, n, n)), symmetry=(1, 2))) raises ValueError: out symmetry does not match result symmetry
control: exp(S, out=empty((n, n))) billed 4,352; exp(S, out=as_symmetric(zeros((n, n)), symmetry=(0, 1))) billed 4,352, result tag SymmetryGroup([1, 0], axes=(0, 1))
control: exp(broadcast_to(S, (3, n, n))) billed 13,056, result tag SymmetryGroup([1, 0], axes=(1, 2)), identical True
control: exp(S, out=zeros((3, n, n))) billed 8,192 (the zeros carried the inferred tag SymmetryGroup([1, 0], axes=(1, 2))), identical True
Expected
These calls compute what numpy computes. The operand's group is first moved into the frame of out (broadcast_group(symmetry, input_shape=..., output_shape=out.shape), which _pointwise_symmetry already uses for operands of different rank), and that aligned group is what the result is checked against and what the tag of out= is compared with, so the last case keeps the (1, 2) tag as exp(broadcast_to(S, (3, n, n))) does.
Where
_counted_unary(src/flopscope/_pointwise.py:1104):symmetry = _symmetry_of(x)(line 1119) is the operand's group; it goes unchanged into_prepare_symmetric_out(out, symmetry)(line 1126) and_wrap_result(result, out=out, symmetry=symmetry)(line 1294)._counted_binary(line 1518):_pointwise_symmetryaligns the operands tonp.broadcast_shapes(x_view.shape, y_view.shape)(line 1548), which leavesoutaside; with one scalar operand the other operand's group is used as is (line 1552)._wrap_result(line 989): with an untaggedoutit calls_validate_result_symmetry(result, symmetry)(line 992) on the wholeoutbuffer; with a taggedout,_prepare_symmetric_out(line 491) compares the tag with the unaligned group and raises (line 503).matmulgoes through_einsum_routed_binary(line 5528), which validates the batchedoutbuffer against the group of the 2-D product (line 5669).
Suggested fix
When out is given and its shape differs from the operands' broadcast shape, align the group to out.shape once, right after it is computed (in _counted_unary, in _counted_binary through the output_shape given to _pointwise_symmetry and in its scalar branch, and in _einsum_routed_binary), and use the aligned group for _prepare_symmetric_out, for _wrap_result and for the tag of the returned out.
Impact
Correctness: valid numpy code raises on both stacks whenever a symmetric-tagged operand (from as_symmetric, or any result that inherits a tag, such as exp(S), S + S or S @ S) is written into a preallocated batch buffer. The workarounds (an fnp.zeros buffer, or broadcast_to first) depend on an inferred tag or on writing the code differently from numpy.
Activity