[Bug] Cannot partition the same array on multiple dimensions

[zzzDavid]

Describe the bug

Schedule.partition() raises AlloValueError("Cannot partition the same array twice") when attempting to partition the same array on two different dimensions. This prevents valid hardware designs where a multi-dimensional array needs distinct partitioning strategies per dimension (e.g., a 3D weight tensor in a systolic array design).

The restriction is an artificial Python-level check in allo/customize.py (lines 309–317). The MLIR backend (mlir/lib/Transforms/LoopTransformations.cpp) already supports composing multiple partition layout maps on the same array.

Reproduction

import allo
from allo.ir.types import int32

def multi_dim_partition(A: int32[8, 8, 8]) -> int32[8, 8, 8]:
    B: int32[8, 8, 8] = 0
    for i, j, k in allo.grid(8, 8, 8):
        B[i, j, k] = A[i, j, k]
    return B

s = allo.customize(multi_dim_partition)
s.partition(s.A, dim=1, factor=4)
s.partition(s.A, dim=2, factor=4)  # raises AlloValueError

Buggy output

allo.backend.AlloValueError: Cannot partition the same array twice: multi_dim_partition:A

Expected behavior

Both partition() calls should succeed, applying block/cyclic partitioning independently on dimensions 1 and 2 of the same array. The MLIR backend already handles layout map composition for this case — the Python-level check should be relaxed to allow partitioning the same array on different dimensions.

Additional context

This issue directly affects the FEATHER accelerator design (examples/feather/feather.py), which currently works around the limitation by using a full partition on the weights array instead of selectively partitioning two dimensions:

# Desired:
# s.partition("top:weights", dim=2, factor=AW)
# s.partition("top:weights", dim=3, factor=AH)

# Current workaround (suboptimal — fully partitions the entire array):
s.partition("top:weights", dim=0)