Simplify PartitionedCall to PartitionedCallOp

PiperOrigin-RevId: 521546933

Author: tf-transform-team

tensorflow_transform/graph_tools_test.py

@@ -25,9 +25,7 @@
 from tfx_bsl.types import tfx_namedtuple
 
 # pylint: disable=g-direct-tensorflow-import
-from tensorflow.python.framework import function
 from tensorflow.python.ops import control_flow_ops
-from tensorflow.python.ops import functional_ops
 # pylint: disable=g-enable-tensorflow-import
 mock = tf.compat.v1.test.mock
 
@@ -60,22 +58,6 @@ def foo(x, y):
   return {'x': x, 'y': y, 'z': a + b + c, 'r': a, 'q': foo(x, x)[0]}
 
 
-def _create_graph_with_tf_defun():
-  x = tf.compat.v1.placeholder(tf.int64)
-  y = tf.compat.v1.placeholder(tf.int64)
-
-  @tf.function
-  def identity(x):
-    return x
-
-  @function.Defun(tf.int64, tf.int64)
-  def foo(x, y):
-    return identity(x * 2), identity(x + y), identity(y * 2)
-
-  a, b, c = functional_ops.partitioned_call([x, y], foo)
-  return {'x': x, 'y': y, 'z': a + b + c, 'r': a, 'q': foo(x, x)[0]}
-
-
 def _create_graph_with_tf2_saved_model_with_table():
 
   def export_saved_model_v2():
@@ -674,12 +656,6 @@ def testInitializableGraphAnalyzerReadyToRunRaises(
           feeds=['x', 'y'],
           fetches=['z'],
           expected_dependent_inputs=['x', 'y']),
-      dict(
-          testcase_name='_tf_defun',
-          create_graph_fn=_create_graph_with_tf_defun,
-          feeds=['x', 'y'],
-          fetches=['z'],
-          expected_dependent_inputs=['x', 'y']),
       dict(
           testcase_name='_tf_function_signature_forces_dependencies',
           create_graph_fn=_create_graph_with_tf_function,