Shaped-array materialization can violate the unbounded-shape invariant

[bdelwood]

Describe the Bug

First off, huge thanks for pushing forward the state-of-the-art on array typing in the Python ecosystem.

I encountered an issue while trying to add array typing to a jax-based analysis codebase. The main numerical interfaces accept properties as scalars or arrays, convert them to jax arrays, and pass them through shape-preserving overloads. For example,

from typing import Any, overload

from shape_extensions import SizeTuple, shaped_array


@shaped_array(shape="Shape")
class Array[Shape: SizeTuple = tuple[Any, ...]]: ...


# Stand-ins for jax.Array and jax.numpy.asarray.
type PermittivityLike = Array | float


def asarray(value: PermittivityLike) -> Array: ...


# Reduced form of JAX's shape-preserving dtype conversion overload.
@overload
def as_complex[S: SizeTuple](value: Array[S]) -> Array[S]: ...


@overload
def as_complex(value: PermittivityLike) -> Array: ...


def as_complex(value: PermittivityLike) -> Array: ...


def assemble_permittivity(permittivity: PermittivityLike) -> Array:
    return as_complex(asarray(permittivity))

panics with

shaped-array unbounded shapes must be tuple[Any, ...]

I think that the problem is that materialization changes the Any inside a gradual shape without running the usual shape normalization afterward. That leaves ShapedArrayShape in a form it normally does not allow. The panic only happens later, when display or tuple projection assumes the invariant still holds.

Happy to discuss a PR, but I am fairly new to Rust and have never dived this deep into Python typing.

Sandbox Link

https://pyrefly.org/sandbox/?project=N4IgZglgNgpgziAXKOBDAdgEwEYHsAeAdAA4CeS4ATrgLYAEALqcROgOZ0Q3G6UN0BBdKQA0dXADcYlKLlSYAOuiVhq9OAAtUxGAH0Y+BjHRwIuE5269+AZQgAvGABUArsVhjN2mJl2pKlKikSiHoAAJeOr7+gaQAFJEwALwKIDZaOqkAlEoAxlCocHCCAUEA2uneiHR2jq7uMHRJjG6wZUKidITdALo91d2EoUoAxDUMGJgAtKzFYLx0AFaoRAKlpHSTSyuE6C7cpISFMUFD6Ew6dAAK0jQQDAwQEvekADIQANaNzWuxdAA+dDAslQDGGWBgYE2aHWcQkqCgLhg1RulDuDyeL3eXyydCmAD4SrEBt1wWMAEo+Fy5HxA3j0XBQgBSAgAGgByYqJKbESjwaTPdh0TAXRq5cxSSimcziSUgzBnMKSaTypSYSHQ3Ti7iwfAVaq1ZytGA9OEIpHVX7lGw9XEEonW-pdUnKV1KuVyRQQqGFLW0Br4M2I5HXW73R7PJjYmB2wlW0gks6hdU+uB+nUGIMW0No8OYqOfGN4uPrRPglPQuAwGjYWC6HS5jGR+IN9ERl4osNNrGF2MOhNKOhDuh8hguSjoTXagNxY6w1t55tZHLoEAiECJfSGYzSkwkUgQCiqWiMZisDhcHh8QTCUKJaKlm-BV35QrFQ31WCIQfDiu6LUFEUuhsDAjxGDQ-5xPkcBiMQ-ioDQcCWsIfYdGWq7rmQfLAocDC0FAFBkAwGjmFMkq7k0dCpAAzIQACM1GpEoIAAL7rqguQRjAABi0AwBQaBYHgRBkKxQA

(Only applicable for extension issues) IDE Information

No response