pycute: composition returns a layout that is not A ∘ B when A does not distribute over B's modes #6

Version: NVlabs/CuTe 111253d17e2f0f8631f43999b43ac4afa5954b04 (pycute 0.1.0).

Reproducer

from pycute import *
A = Layout((4, 8), (8, 1))
B = Layout((2, 3), (3, 1))
R = composition(A, B)
print(R)                                   # (2, 3):(24, 8)
print([R(i) for i in range(size(B))])      # [0, 24, 8, 32, 16, 40]
print([A(B(i)) for i in range(size(B))])   # [0, 24, 8, 1, 16, 9]

R(3) = 32, but A(B(3)) = A(4) = 1. The same happens at i = 5. composition raises nothing, although its docstring promises result(i) == A(B(i)) for i in range(size(B)), which is whitepaper Eq. (17).

Cause

For a B whose shape is a tuple, Layout._composition distributes over B's modes, A ∘ (B0, B1, ...) = (A ∘ B0, A ∘ B1, ...) (whitepaper Eq. 22). It composes each mode separately and concatenates the results. Here A ∘ 2:3 = 2:24 and A ∘ 3:1 = 3:8 are each correct. Their concatenation is not A ∘ B, because A's shape does not distribute over the sum B0(c0) + B1(c1). For example, A(3 + 1) = 1, while A(3) + A(1) = 24 + 8 = 32.

The whitepaper (arXiv:2603.02298, §3.3.2, "Reductive Case: Distributive") makes Eq. (22) conditional on Eq. (23), the distributivity of S over the sublayouts of B. It gives two sufficient conditions: every base stride d_i of B meets the stride divisibility condition (20), and the base sublayouts s_i : d_i have mutually segregated images (s_i d_i ≤ d_j or s_j d_j ≤ d_i). In the example the images of 2:3 ({0, 3}) and 3:1 ({0, 1, 2}) interleave, so the second condition fails. pycute checks neither condition.

Consequences

The error propagates to everything built on composition: logical_divide, logical_product, zipped_divide, blocked_product and raked_product. For example, logical_product(Layout((2,2),(5,10)), Layout((2,2),(1,4))) returns ((2,2),(2,2)):((5,10),(1,4)), which is not injective (offsets 5, 10 and 15 repeat), although both operands are injective. In a corpus of about 2,000 static compositions with small right-hand sides, 69 hit this, and the divides and products built on them add 43 more.

Suggested fix

In Layout._composition, in the is_tuple(B.shape) branch, check Eq. (23) before returning the concatenation. Either:

• decide it exactly for static layouts, by checking R(c) == A(B(c)) for c in range(size(B)) (or the equivalent A(sum_k c_k d_k) == sum_k A(c_k d_k)) and raising ValueError("Distributivity condition violated: ...") otherwise; or, at least,
• check the whitepaper's two sufficient conditions (stride divisibility of every base stride against A's prefix products, and mutually segregated images of B's base modes) and raise when they fail. This may reject some compositions that happen to distribute.

A test such as composition(Layout((4,8),(8,1)), Layout((2,3),(3,1))) raising ValueError would cover it.