facebook / facebook/pyrefly

Shaped-array materialization can violate the unbounded-shape invariant

Open
#4,129 2 comments 0 reactions 1 assignee Claimed by @stroxler View on GitHub
needs-triage stale tensor-shapes typechecking
Dominant language
Rust
Stars
7k
Forks
516
PR merge metrics
No merged PRs in 30d

Description

### 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,

```python
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

```text
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_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.