odmltorch tf integration cleanup

PiperOrigin-RevId: 711819517

Author: chunnienc

ai_edge_torch/odml_torch/export.py

@@ -222,11 +222,6 @@ def __call__(self, *args):
     # Lazy importing TF when execution is needed.
     return self.tf_function(*args)
 
-  def to_flatbuffer(self):
-    from . import tf_integration
-
-    return tf_integration.mlir_to_flatbuffer(self)
-
 
 # TODO(b/331481564) Make this a ai_edge_torch FX pass.
 def _convert_i64_to_i32(exported_program: torch.export.ExportedProgram):

ai_edge_torch/odml_torch/test/test_tf_integration.py

@@ -44,29 +44,6 @@ def forward(self, x, y):
 
     self.assertTrue(np.allclose(lowering_output, torch_output))
 
-  @googletest.skip("b/353280409")
-  def test_dynamic_mlir_lowered_call(self):
-    class AddModel(torch.nn.Module):
-
-      def forward(self, x, y):
-        return x + y + x + y
-
-    model = AddModel().eval()
-    batch = torch.export.Dim("batch")
-    ep = torch.export.export(
-        model,
-        (torch.rand((2, 10)), torch.rand((2, 10))),
-        dynamic_shapes={"x": {0: batch}, "y": {0: batch}},
-    )
-
-    lowered = odml_torch.export.exported_program_to_mlir(ep)
-
-    val_args = (torch.rand((10, 10)), torch.rand((10, 10)))
-    torch_output = model(*val_args).detach().numpy()
-    lowering_output = np.array(lowered(*val_args))
-
-    self.assertTrue(np.allclose(lowering_output, torch_output))
-
   def test_resnet18(self):
     model = torchvision.models.resnet18().eval()
     forward_args = lambda: (torch.rand((1, 3, 224, 224)),)

ai_edge_torch/odml_torch/tf_integration.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""APIs to convert lowered MLIR from PyTorch to TensorFlow and TFLite artifacts."""
+"""APIs to convert lowered MLIR from PyTorch to TensorFlow artifacts."""
 
 import re
-import tempfile
 
 import tensorflow as tf
 import torch
@@ -155,46 +154,3 @@ def mlir_to_tf_function(lowered: export.MlirLowered):
       _wrap_as_tf_func(lowered, tf_state_dict),
       input_signature=_make_input_signatures(lowered),
   )
-
-
-def mlir_to_flatbuffer(lowered: export.MlirLowered):
-  """Convert the MLIR lowered to a TFLite flatbuffer binary."""
-  tf_state_dict = _build_tf_state_dict(lowered)
-  signature_names = [tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
-  tf_signatures = [_make_input_signatures(lowered)]
-  tf_functions = [_wrap_as_tf_func(lowered, tf_state_dict)]
-
-  tf_module = tf.Module()
-  tf_module.f = []
-
-  for tf_sig, func in zip(tf_signatures, tf_functions):
-    tf_module.f.append(
-        tf.function(
-            func,
-            input_signature=tf_sig,
-        )
-    )
-
-    tf_module._variables = list(tf_state_dict.values())
-
-  tf_concrete_funcs = [
-      func.get_concrete_function(*tf_sig)
-      for func, tf_sig in zip(tf_module.f, tf_signatures)
-  ]
-
-  # We need to temporarily save since TFLite's from_concrete_functions does not
-  # allow providing names for each of the concrete functions.
-  with tempfile.TemporaryDirectory() as temp_dir_path:
-    tf.saved_model.save(
-        tf_module,
-        temp_dir_path,
-        signatures={
-            sig_name: tf_concrete_funcs[idx]
-            for idx, sig_name in enumerate(signature_names)
-        },
-    )
-
-    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
-    tflite_model = converter.convert()
-
-  return tflite_model