Add support for Python's optimized mode to torchscript

[luisasanmartin]

🚀 Feature

Add support for Python's optimized mode to torchscript.

Motivation

When debugging a model, it’s useful to add asserts and other similar logic to catch bugs. However, these asserts can be slow, so for production use cases, it’s useful to remove them all. Running with python -O will do that. However, when exporting a model, the asserts remain in the compiled graph even with the optimized flag turned on.

Pitch

Be able to export a model with python -O and have it remove all asserts and code guarded with if debug

[ansley]

cc @gmagogsfm Is this something that we'd be able to work on? It seems like it would be a larger project

[suo]

One consideration is that often times asserts are semantically important for type safety. For example, if you use an assert to refine an Optional[Tensor] to Tensor, we need to preserve that assertion so that the resulting program stays type safe.

So if we did strip asserts, we'd have to ensure that there were no downstream dependencies on the assertion.

[gmagogsfm]

One consideration is that often times asserts are semantically important for type safety. For example, if you use an assert to refine an Optional[Tensor] to Tensor, we need to preserve that assertion so that the resulting program stays type safe.

So if we did strip asserts, we'd have to ensure that there were no downstream dependencies on the assertion.

Yeah, good point. Because assert is so commonly used to refine shapes, stripping assert is effectively a "lossy" (or even risky) transformation. TorchScript tends to lean on the conservative side when making transformations. What do you think about this option though: We can open up the ability to register custom transformation passes so that you can write a pass that operates on TorchScript IR to remove asserts.