Top level import for compiler optimization passes.

examples/intro_with_auto_param.ipynb

@@ -109,8 +109,7 @@
     }
    ],
    "source": [
-    "import tf_shell.python.shell_optimizers as shell_optimizers\n",
-    "shell_optimizers.enable_tf_shell_optimizer()\n",
+    "tf_shell.enable_optimization()\n",
     "\n",
     "a = [1, 2, 3]\n",
     "b = [4, 5, 6]\n",

examples/label_dp_sgd.ipynb

@@ -52,7 +52,6 @@
     "import keras\n",
     "import numpy as np\n",
     "import tf_shell\n",
-    "import tf_shell.python.shell_optimizers as shell_optimizers\n",
     "import tf_shell_ml"
    ]
   },
@@ -93,7 +92,7 @@
    "outputs": [],
    "source": [
     "# Turn on the shell optimizer to use autocontext.\n",
-    "shell_optimizers.enable_tf_shell_optimizer()\n",
+    "tf_shell.enable_optimization()\n",
     "use_fast_reduce_sum = True\n",
     "\n",
     "m = tf_shell_ml.TfShellSequential(\n",

tf_shell/__init__.py

@@ -44,3 +44,5 @@
 
 from tf_shell.python.shell_ops import clip_and_noise_features_party
 from tf_shell.python.shell_ops import clip_and_noise_labels_party
+
+from tf_shell.python.shell_optimizers import enable_optimization

tf_shell/python/shell_optimizers.py

@@ -139,7 +139,7 @@ def optimize_shell_graph(func, optimizers=all_shell_optimizers):
 
 # Here is a method to enable custom optimizers described by
 # https://github.com/tensorflow/tensorflow/issues/55451#issuecomment-1147065792
-def enable_tf_shell_optimizer(optimizers=all_shell_optimizers):
+def enable_optimization(optimizers=all_shell_optimizers):
     rewriter_config = rewriter_config_pb2.RewriterConfig()
     rewriter_config.meta_optimizer_iterations = rewriter_config_pb2.RewriterConfig.ONE
     for optimizer in optimizers:

tf_shell/test/auto_param_optimizer_test.py

@@ -328,7 +328,7 @@ def test_func(self):
             maxval=127,
         )
 
-        shell_optimizers.enable_tf_shell_optimizer()
+        tf_shell.enable_optimization()
         c = ct_ct_add(a, b, True)
 
         # If the optimizer ran, the shape should be padded out to the

tf_shell/test/ct_pt_optimizer_test.py

@@ -221,7 +221,7 @@ def test_auto_optimize(self):
         # Turn on automatic optimization. Note there is no way to get the
         # optimized graph from the tf.function so we need to rely on timing info
         # to make sure it's turned on.
-        shell_optimizers.enable_tf_shell_optimizer(["CtPtOptimizer"])
+        tf_shell.enable_optimization(["CtPtOptimizer"])
 
         opt_time = timeit(lambda: ct_pt_pt_add(ct_a, pt_b, 10), number=10)
 

tf_shell/test/pt_pt_optimizer_test.py

@@ -280,7 +280,7 @@ def test_auto_optimize(self):
         # Turn on automatic optimization. Note there is no way to get the
         # optimized graph from the tf.function so we need to rely on timing info
         # to make sure it's turned on.
-        shell_optimizers.enable_tf_shell_optimizer(["PtPtOptimizer"])
+        tf_shell.enable_optimization(["PtPtOptimizer"])
 
         opt_time = timeit(
             lambda: pt_add(a, b, 10, test_context.shell_context, test_context.key),

tf_shell_ml/test/model_distrib_test.py

@@ -84,7 +84,7 @@ def test_model(self):
         val_dataset = val_dataset.batch(batch_size, drop_remainder=True)
 
         # Turn on the shell optimizer to use autocontext.
-        shell_optimizers.enable_tf_shell_optimizer()
+        tf_shell.enable_optimization()
 
         m = tf_shell_ml.TfShellSequential(
             [

tf_shell_ml/test/model_local_test.py

@@ -40,7 +40,7 @@
 class TestModel(tf.test.TestCase):
     def test_model(self):
         # Turn on the shell optimizer to use autocontext.
-        shell_optimizers.enable_tf_shell_optimizer()
+        tf_shell.enable_optimization()
 
         m = tf_shell_ml.TfShellSequential(
             [