Shaped-array materialization can violate the unbounded-shape invariant
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
(Only applicable for extension issues) IDE Information
No response
Source: facebook/pyrefly