Namespace cleanup (#229)

Two main things happened:
- I removed {wrap_key, PythonTensor, pythonkey_trace} from being public
APIs
- I moved all compilation related things to the functorch.compile
namespace. This includes nnc_jit which is now in
functorch.compile.nnc_jit

Concerns:
- nnc_jit was in the functorch namespace for a long time. Should we
leave it there? Are there stakeholders to notify?

Author: zou3519

benchmarks/operator_authoring.py

@@ -3,7 +3,7 @@
 import pandas as pd
 import timeit
 import torch
-from functorch import pointwise_operator
+from functorch.compile import pointwise_operator
 
 WRITE_CSV = False
 CUDA = False

functorch/__init__.py

@@ -9,19 +9,24 @@
 import textwrap
 from . import _C
 
+# Top-level APIs. Please think carefully before adding something to the
+# top-level namespace:
+# - private helper functions should go into functorch._src
+# - very experimental things should go into functorch.experimental
+# - compilation related things should go into functorch.compile
+
+# functorch transforms
 from ._src.vmap import vmap
 from ._src.eager_transforms import grad, grad_and_value, vjp, jacrev
+from ._src.python_key import make_fx
+
+# utilities. Maybe these should go in their own namespace in the future?
 from ._src.make_functional import (
     make_functional_with_buffers,
     make_functional,
     combine_state_for_ensemble,
     FunctionalModule,
 )
-from ._src.python_key import wrap_key, PythonTensor, pythonkey_trace, make_fx, nnc_jit, make_nnc
-from ._src.nnc_compile import nnc_compile, get_ops
-from ._src.eager_compilation import compiled_function, compiled_module, tvm_compile, draw_joint_graph, default_partition, partition_with_recompute_fwd_in_bwd
-from ._src.operator_authoring import pointwise_operator
-
 
 # Monkeypatching lol
 _old_cross_entropy = torch.nn.functional.cross_entropy

functorch/compile/__init__.py

@@ -1 +1,11 @@
-from .._src.operator_authoring import pointwise_operator
+from .._src.operator_authoring import pointwise_operator
+from .._src.python_key import nnc_jit, make_nnc
+from .._src.nnc_compile import nnc_compile, get_ops
+from .._src.eager_compilation import (
+    compiled_function,
+    compiled_module,
+    tvm_compile,
+    draw_joint_graph,
+    default_partition,
+    partition_with_recompute_fwd_in_bwd,
+)

test/test_pythonkey.py

@@ -23,7 +23,10 @@
 import functorch
 from functorch import (
     grad, vjp, vmap, jacrev, grad_and_value,
-    make_fx, nnc_jit, compiled_function, compiled_module,
+    make_fx,
+)
+from functorch.compile import (
+    nnc_jit, compiled_function, compiled_module,
     partition_with_recompute_fwd_in_bwd
 )