This is conceptually very simple but adds code complexity. It would work automatically for functions under syntax extensions. We would apply the function to a "tracing tensor" and would store the inferred shape constraints (not shapes since we don't know what the function returns, it could be more than just a single tensor). Application would then freshen the constraint and include it with equalities for the correspondence between the tracing tensor and applied tensor shapes.
The two big benefits are: earlier shape error catching -- although some errors still delayed to "finish inference" passes; and improvement to shape inference efficiency (with a downstream benefit to shape inference debugging).
This is conceptually very simple but adds code complexity. It would work automatically for functions under syntax extensions. We would apply the function to a "tracing tensor" and would store the inferred shape constraints (not shapes since we don't know what the function returns, it could be more than just a single tensor). Application would then freshen the constraint and include it with equalities for the correspondence between the tracing tensor and applied tensor shapes.
The two big benefits are: earlier shape error catching -- although some errors still delayed to "finish inference" passes; and improvement to shape inference efficiency (with a downstream benefit to shape inference debugging).