[Functorch] Don't try to setup docs if called with optimize

Test plan: `python3 -OO -c "import torch._functorch.deprecated"`

torch/_functorch/deprecated.py

@@ -7,6 +7,7 @@
 import textwrap
 from typing import Any, Callable, Optional, Tuple, Union
 import warnings
+import sys
 
 """
 The APIs in this file are exposed as `functorch.*`. They are thin wrappers
@@ -105,18 +106,20 @@ def combine_state_for_ensemble(models):
     warn_deprecated('combine_state_for_ensemble', 'torch.func.stack_module_state')
     return _nn_impl.combine_state_for_ensemble(models)
 
-setup_docs(vmap, _vmap_impl.vmap, 'torch.vmap')
-setup_docs(grad)
-setup_docs(grad_and_value)
-setup_docs(vjp)
-setup_docs(jvp)
-setup_docs(jacrev)
-setup_docs(jacfwd)
-setup_docs(hessian)
-setup_docs(functionalize)
-setup_docs(make_functional, _nn_impl.make_functional,
-           'torch.func.functional_call')
-setup_docs(make_functional_with_buffers, _nn_impl.make_functional,
-           'torch.func.functional_call')
-setup_docs(combine_state_for_ensemble, _nn_impl.combine_state_for_ensemble,
-           'torch.func.stack_module_state')
+if sys.flags.optimize < 2:
+    setup_docs(vmap, _vmap_impl.vmap, 'torch.vmap')
+    setup_docs(grad)
+    setup_docs(grad_and_value)
+    setup_docs(vjp)
+    setup_docs(jvp)
+    setup_docs(jacrev)
+    setup_docs(jacfwd)
+    setup_docs(hessian)
+    setup_docs(functionalize)
+    setup_docs(make_functional, _nn_impl.make_functional,
+               'torch.func.functional_call')
+    setup_docs(make_functional_with_buffers, _nn_impl.make_functional,
+               'torch.func.functional_call')
+    setup_docs(combine_state_for_ensemble, _nn_impl.combine_state_for_ensemble,
+               'torch.func.stack_module_state')
+