Record shaping assertions and use them for tracing / scripting optimization and codegen

[vadimkantorov]

Motivated in #40373 (comment).

Recording assertions assert X.shape[1] == Y.shape[2] and X.shape[2] == 64 and Z.ndim == 3 (or directly via torch.assert_shapes(...)) as part of tracing / scripting can help and streamline codegen (as can eliminate unneeded checks or broadcasts). #40373 contains a wider discussion, but @ezyang suggested that a more limited scope can be of value as well.

cc @gmagogsfm @apaszke

[vadimkantorov]

If these are recorded, it may be useful to let user access these (shape inference) "type hints".
One way this could be done is by: assert torch.record_shapes(lambda X, Y, Z: X.shape[1] == Y.shape[2] and X.shape[2] == 64 and Z.ndim == 3, X, Y, Z) and then accesssing by hints = torch.get_recorded_shape_hints(X, Y, Z)

[vadimkantorov]

Btw with dynamo, this idea of explicit user-provided hints/guarantees may be still relevant @ezyang @albanD (including for asserting numel less than int32 max)

Some other useful constraints can include divisibility: e.g. assert tensor.shape[-1] % 64 == 0 which could simplify the emitted code

If these are in form of assert statement, they can be nicely executed in eager mode as well and blow up nicely if something is wrong

[vadimkantorov]

Related: #107735

[vadimkantorov]

Related idea on manual decaying shapes to static to help/simplify the export:

• aten::squeeze exported to ONNX as an If node #109292

[vadimkantorov]

Somewhat related question on Dynamo preserving the shape/dtype assert statements in Dynamo-captured graphs: https://stackoverflow.com/questions/78426965/how-do-i-get-dynamo-export-to-not-export-shape-asserts