def METHOD_NAME(comp):
  """Counts TF Variables in `comp` if `comp` is a TF block."""
  py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock)
  if (
      not isinstance(comp, building_blocks.CompiledComputation)
      or comp.proto.WhichOneof('computation') != 'tensorflow'
  ):
    raise ValueError(
        'Please pass a '
        '`building_blocks.CompiledComputation` of the '
        '`tensorflow` variety to `count_tensorflow_variables_in`.'
    )
  graph_def = serialization_utils.unpack_graph_def(
      comp.proto.tensorflow.graph_def
  )
  def _node_is_variable(node):
    # TODO: b/137887596 - Follow up on ways to count Variables on the GraphDef
    # level.
    op_name = str(node.op).lower()
    return (
        op_name.startswith('variable') and op_name not in ['variableshape']
    ) or op_name == 'varhandleop'
  def _count_vars_in_function_lib(func_library):
    total_nodes = 0
    for graph_func in func_library.function:
      total_nodes += sum(
          _node_is_variable(node) for node in graph_func.node_def
      )
    return total_nodes
  return sum(
      _node_is_variable(node) for node in graph_def.node
  ) + _count_vars_in_function_lib(graph_def.library)