Repository navigation
NumPy shape stubs: Generator.random / standard_normal / standard_exponential / standard_cauchy return unshaped arrays #5595
Description
Describe the Bug
The Generator methods random, standard_normal, standard_exponential and standard_cauchy don't manipulate matrices but only generate random values of a desired shape, so the return shape is just the size parameter if provided, otherwise a float.
import numpy as np
rng = np.random.default_rng(0)
a = np.zeros((10, 10))
b = rng.standard_normal((10, 10))
w = np.zeros((5, 5))
a @ w # pyrefly: error, `@` is not supported between [10, 10] and [5, 5]
b @ w # pyrefly: no error, but crashes at runtime with a matmul ValueErrorb is inferred as ndarray[tuple[Any, ...], dtype[float64]]. The same happens with rng.random, rng.standard_exponential and rng.standard_cauchy.
Expected: the result shape follows size, as with np.zeros, so rng.standard_normal((10, 10)) is ndarray[[10, 10], dtype[float64]]. With no size, the result stays a float.
I have prepared 2 PRs to add shapes to these functions, since I use them daily:
rng.randomandrng.standard_cauchyrng.standard_normalandrng.standard_exponential
Sandbox Link
No response
(Only applicable for extension issues) IDE Information
No response
Activity
AlessandroRuggiero commentedon Nov 9, 2026
#claim
github-actions commentedon Nov 9, 2026
@AlessandroRuggiero you've claimed this issue — it's now assigned to you. Thanks for picking it up! 🎉
stroxler commentedon Nov 9, 2026
Thanks!
Unrelated to the bug / PRs - sounds like you are using Pyrefly with numpy and a lot of random functionality. How's it working for you overall so far?
This feature is still early, I'm starting to roll it out internally at Meta but mostly on torch code. I'm very curious to hear from numpy users, especially anyone doing statistical / simulation work since I used to spend a lot of time on that and I'm trying to build something that would help my former (and perhaps future) self.
I'm aware the numpy stubs are less complete, partly a consequence of our having tested torch more against real codebases; I find new gaps every time I annotate a real codebase. The improvements are much appreciated
Infer NumPy Generator.random and standard_cauchy result shapes (#5596)
AlessandroRuggiero commentedon Nov 10, 2026
Thanks, happy to help!
This project is something I have been searching for for a long time; I always wanted a way to use generics-like syntax to handle matrix dimensions.
I noticed that keeping track of the dimensions in my code took up a big portion of my time, and it always felt like a problem other areas of programming had already solved with generics. Rust's const generics are a good example, where an array's length is part of its type and the compiler checks it.
Once my code is annotated, I almost never have to scatter
print(x.shape)calls throughout my codebase. However, I found that the NumPy stubs are still missing a lot of functions. I also briefly tried annotating a PyTorch codebase, and there it worked a lot better, so I realised the problem was mostly missing stubs and decided to add them myself.Adding stubs was pretty easy; the main difficulty is reading NumPy's docs thoroughly, because its functions have lots of overloads, some of which I never knew existed.
I'll keep adding stubs as I run into gaps. Next, I'm planning to cover more of the random
Generator(starting withnormal) and to makeravel/flattenkeep the array's size instead of returning an unknown length.AlessandroRuggiero commentedon Nov 10, 2026
As mentioned, here's the
ravel/flattenchange: #5604. It turned out to be small, since theflatten_shaperule already existed.Also, were the sizes of the PRs good so far or is there something I should do differently in the future?
stroxler commentedon Nov 10, 2026
I think the sizes of the PRs are good, these have been helpful! I think it's probably optional to file an issue for stub-only bugs / gaps in coverage. It's fine to just send PRs in, and probably a good idea to assign them to me.
If you find a situation where Pyrefly itself doesn't work (bad error messages, the dsl is missing functionality, or an outright bug), that calls for an issue I think. You're welcome to attempt a PR for that kind of problem too, but it's more likely to involve some back-and-forth if we're changing the rust code.