alt Hacker NewsThe Pyrefly type checker is starting to work on this kind of shape hinting - so far it only works on Torch but I believe the plan is for it to work with other array packages (eg. JAX, NumPy)
Replies
mathisfun123 • yesterday at 6:17 PM
Shape functions and shape analysis are basically mundane infra in almost every ML compiler/language/DSL.
➕ show 1 reply
See also jaxtyping which, contrary to what its name might imply, covers JAX/PyTorch/NumPy/MLX/TensorFlow arrays and tensors.
https://docs.kidger.site/jaxtyping/