extension: N > 2 differential testing shape logic

[clay-arras]

Algorithm

Function Signature
def common_reshape(a_shape: list[int], b_shape: list[int]) -> tuple[list[int], list[int]]
This spec defines a helper that finds common shapes using reshape and expand_dims. This algorithm is slightly modified from the original issue #15.

Movement: Right to Left
Assume $x$ is the current shape value at index $i$ for tensor A, and $y$ is the current shape value at index $i$ for tensor B.

• Replace $x$ with $gcd(x, y)$
• Replace $y$ with $gcd(x, y)$

If there are no more indices on the left, pad with ones:

• Carry over $x / gcd(x, y)$ to the nearest next index on the left.
• Carry over $y / gcd(x, y)$ to the nearest next index on the left.

Edge Case
Finally, if BOTH A and B have no more indices on the left, AND $x$ and $y$ are coprime ($gcd(x, y) == 1$), fall back on this:

• [x, 1, ...]
• [1, y, ...]

---

Step-by-Step Trace

   [4, 6]
[5, 2, 3]
       ^

   [8, 3]
[5, 2, 3]
    ^    

[1, 8, 3]
[5, 2, 3]
    ^    

[4, 2, 3]
[5, 2, 3]
 ^       

[4, 1, 2, 3]
[1, 5, 2, 3]
    ^       

---

Examples

Example 1

• Input A: [3, 4, 1, 3]
• Input B: [5, 2, 5]

Expected Result:

• A: [3, 1, 4, 1, 3]
• B: [1, 5, 1, 10, 1]

Example 2

• Input A: [7, 1, 3]
• Input B: [2, 6]

Expected Result:

• A: [7, 1, 3]
• B: [4, 3]

---

Motivation

Allows us to perform a binary operation on ANY two tensors, regardless of shape, and let the broadcasting be a valid operation.

---

Note: Matmul

If the operation is matmul, we do [a / gcd(a, b), gcd(a, b)] and [gcd(a, b), b / gcd(a, b)] where $a$ and $b$ are the product of the first tensor's shape and the second tensor's shape, respectively.

Example

• Input A: [7, 9]
• Input B: [3, 4, 1]

Expected Result:

• A: [21, 3]
• B: [3, 4]