def _get_table_init_op_and_inputs(self, table_init_op_or_tensor):
    """Get a tuple of table init op and keys for its input ops."""
    # If a TF2 exported SavedModel with a table is loaded inside the
    # preprocessing_fn, the TABLE_INITIALIZERS collection of the outer graph
    # contains a Tensor whose parent op is of type StatefulPartitionedCall.
    # The nested func graph for this StatefulPartitionedCall contains the
    # table initializer.
    if (isinstance(table_init_op_or_tensor, tf.Tensor) and
        table_init_op_or_tensor.op.type == 'StatefulPartitionedCall'):
      result = (table_init_op_or_tensor.op,
                [input_t.op for input_t in table_init_op_or_tensor.op.inputs])
    else:
      assert isinstance(table_init_op_or_tensor, tf.Operation)
      # We are using the table init op information and the table op information,
      # since that is a unique description of the table op.
      table_ops = []
      for input_t in table_init_op_or_tensor.inputs:
        # One of the inputs to the initializer op should be the table op. If
        # no table op is found, (as in the case of a StatefulPartitionedCall)
        # all inputs are added to the source dict.
        if input_t.dtype == tf.resource:
          table_ops.append(input_t.op)
      assert len(table_ops) == 1
      result = (table_init_op_or_tensor, [table_ops[0]])
    return result
Get a tuple of table init op and keys for its input ops.
tensorflow_transform/graph_tools.py