func::FuncOp GetEntryFuncOp(TF::XlaCallModuleOp xla_call_module_op,
                            const SymbolTable symbol_table) {
  const auto entry_function_symbol_ref =
      xla_call_module_op->getAttrOfType<FlatSymbolRefAttr>(kEntryFuncAttrName);

  return dyn_cast_or_null<func::FuncOp>(
      symbol_table.lookup(entry_function_symbol_ref.getValue()));
}

      llvm::dyn_cast<tf_device::LaunchOp>(preprocess_op->getParentOp());
  // Device of original launch looked up before moving the preprocessing ops
  // around.
  StringAttr device =
      original_launch_op->getAttrOfType<StringAttr>(kDeviceAttr);

  if (!device) {
    return original_launch_op->emitOpError()
           << "Launch op has an invalid device attribute.";
  }

  // Find the TPUCompile successor.

           << "]";

  return success();
}

bool VariableIsInitialized(TF::VarHandleOp var_handle_op) {
  auto is_variable_initialized =
      var_handle_op->getAttrOfType<BoolAttr>("_is_initialized");
  // Assume variable is initialized if attribute is not set.
  // There are paths that doesn't mark the variables. All variables
  // that doesn't have the attribute will be promoted.

  // The StringRef cast is necessary before cxx14.
  mlir::IntegerAttr group_key_attr =
      module_op->getAttrOfType<mlir::IntegerAttr>(
          mlir::StringRef(kGroupKeyAttrName.data(), kGroupKeyAttrName.size()));
  mlir::IntegerAttr group_size_attr =
      module_op->getAttrOfType<mlir::IntegerAttr>(mlir::StringRef(
          kGroupSizeAttrName.data(), kGroupSizeAttrName.size()));

  LogicalResult matchAndRewrite(Operation* op,
                                PatternRewriter& rewriter) const override {
    auto tflite_indices_attr =
        op->getAttrOfType<ArrayAttr>(kTfLiteInputIndices);
    if (!tflite_indices_attr) return failure();

    SmallVector<int64_t, 20> tflite_indices;
    for (auto index_attr : tflite_indices_attr.getValue()) {

// this pass. This pass ignores _xla_map_outside_compilation, which will only be
// handled by extract_outside_compilation pass.
bool HasOutsideCompilationAttribute(Operation* op) {
  return op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr) != nullptr &&
         !op->hasAttrOfType<BoolAttr>(kXlaMapOutsideCompilationAttr);
}

// Finds op that created a given value. If the value is a BlockArgument, this

                         "composite_dot_general_fn", operands, results,
                         attributes)[0]
          .getDefiningOp();
  const auto entry_function_symbol_ref =
      lifted_op->getAttrOfType<FlatSymbolRefAttr>("_entry_function");
  SymbolTable symbol_table(*module_op);
  auto entry_func = dyn_cast_or_null<func::FuncOp>(
      symbol_table.lookup(entry_function_symbol_ref.getValue()));

                                         solver)) {
        op_only_operates_on_composite_devices = true;
      }
    }

    auto device_attr = op.getAttrOfType<StringAttr>(kDeviceAttr);
    int64_t device_id;
    if (!device_attr) {
      device_id = kAnyDevice;
    } else if (devices.find(device_attr) != devices.end()) {
      device_id = devices[device_attr];

// only check if it is a host tf.IteratorGetNext.
bool IsSupportedHostInputOp(Operation* op) {
  TF::IteratorGetNextOp iter = llvm::dyn_cast<TF::IteratorGetNextOp>(op);
  if (!iter) return false;
  auto device = op->getAttrOfType<StringAttr>(kDeviceAttr);
  if (!device) return false;
  tensorflow::DeviceNameUtils::ParsedName parsed_device;
  if (!tensorflow::DeviceNameUtils::ParseFullName(device.getValue().str(),

  return true;
}

StringAttr GetReplicationAttr(mlir::Operation* op) {
  return op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr);
}

StringAttr GetReplicationAttr(TF::TPUCompilationResultOp op) {
  // Special case for getting the replication region for
  // TPUCompilationResultsOp.
  return op->getAttrOfType<StringAttr>(kTpuCompilationStatus);
}