flr, function, &fbody, &constant_arg_indices, &resource_arg_indices));
      VLOG(3) << "Found constant arg indices: "
              << absl::StrJoin(constant_arg_indices, ", ");

      node->AddAttr("_input_hostmem", constant_arg_indices);
    }
  }
  return absl::OkStatus();
}

  for (Node* node : graph->nodes()) {
    if (CanCreateXlaKernel(node->def())) {
      EXPECT_FALSE(found);
      found = true;
      std::vector<int32> hostmem_attr;
      EXPECT_TRUE(TryGetNodeAttr(node->def(), "_input_hostmem", &hostmem_attr));
      EXPECT_EQ(hostmem_attr.size(), 1);
      EXPECT_EQ(hostmem_attr[0], 1);
    }
  }
  EXPECT_TRUE(found);
}

}  // namespace