def _lookup_table(table: lookup_ops.LookupInterface, x: tf.Tensor,
                  control_dependency: Optional[tf.Tensor]) -> tf.Tensor:
  """Look up x in table with an optional depndency on control_dependency."""
  with contextlib.ExitStack() as stack:
    # tf.control_dependencies([tensor]) adds a dependency to tensor.op. Wrap the
    # tensor in an identity op to ensure that walking the graph from `result`
    # encounters the control_dependency tensor.
    if control_dependency is not None:
      stack.enter_context(
          tf.control_dependencies([tf.identity(control_dependency)]))
    result = table.lookup(x)
  return result
Look up x in table with an optional depndency on control_dependency.
tensorflow_transform/tf_utils.py