}

mlir::LogicalResult GetDevicesFromOp(mlir::Operation* op,
                                     mlir::TF::RuntimeDevices* devices) {
  auto devices_attr = op->getAttr(kDevicesAttr);
  if (!devices_attr) return mlir::success();

  if (auto array_attr = mlir::dyn_cast<mlir::ArrayAttr>(devices_attr)) {
    return GetDevicesFromOp(op, array_attr, devices);

      return false;
    }
    // FusedMatMul kernel does not support grad_a/grad_b attrs
    if ((matmul->hasAttr("grad_a") &&
         mlir::cast<BoolAttr>(matmul->getAttr("grad_a")).getValue()) ||
        (matmul->hasAttr("grad_b") &&
         mlir::cast<BoolAttr>(matmul->getAttr("grad_b")).getValue())) {
      (void)rewriter.notifyMatchFailure(matmul, [&](Diagnostic &diag) {
        diag << "FusedMatMul kernel does not support grad_a/grad_b attrs";

                       absl::StrCat("module_", module_name));
  auto parent_module = entry_func->getParentOfType<ModuleOp>();
  auto versions_attr = parent_module->getAttr(kVersionsAttr);
  if (!versions_attr)
    return parent_module.emitError(CreateMissingAttributeMsg(kVersionsAttr));

  module_for_func.get().getOperation()->setAttr(kVersionsAttr, versions_attr);

    func::FuncOp ResolveFunc(::mlir::SymbolTableCollection* table) {
      if (table)
        return table->lookupNearestSymbolFrom<func::FuncOp>(*this, getFAttr());
      return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(*this, getFAttr());
    }
    // TODO(b/204997177): Deprecate and remove.
    func::FuncOp func() {  return ResolveFunc(nullptr); }
  }];
}

void GroupByDialectPass::processFunction(
    mlir::func::FuncOp func, int& counter,
    llvm::SmallDenseSet<StringRef>& dialects, Operation* module) {
  // don't re-process functions we generated
  if (func->getAttr("dialect")) return;
  processRegion(func.getBody(), counter, dialects, module);
}

void GroupByDialectPass::processRegion(mlir::Region& region, int& counter,

        rewriter.getDenseI64ArrayAttr(kOihwToHwioPermutation));

    // [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
    const auto new_dimension_nums = rewriter.getAttr<ConvDimensionNumbersAttr>(
        /*inputBatchDimension=*/0, /*inputFeatureDimension=*/3,
        /*inputSpatialDimensions=*/SmallVector<int64_t>{1, 2},
        /*kernelInputFeatureDimension=*/2, /*kernelOutputFeatureDimension=*/3,

    cc_type v;                                                           \
    auto* cc_ctx = reinterpret_cast<InferenceContext*>(ctx);             \
    Status s = cc_ctx->GetAttr(attr_name, &v);                           \
    Set_TF_Status_from_Status(status, s);                                \
    if (s.ok()) {                                                        \

    auto range = llvm::make_filter_range(op->getUsers(), [&](Operation *user) {
      bool same_block = user->getBlock() == op->getBlock();
      bool same_device = op->getAttr(kDeviceAttr) == user->getAttr(kDeviceAttr);
      return same_block && same_device && state.IsMember(user);
    });
    return {range.begin(), range.end()};
  }

  // Users of block argument must be in the same block.

}

// Returns the function attribute for the given call op which is lifted for
// quantization.
inline FlatSymbolRefAttr GetFuncAttr(TF::PartitionedCallOp call_op) {
  return mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
}

inline FlatSymbolRefAttr GetFuncAttr(TF::XlaCallModuleOp call_op) {
  return call_op->getAttrOfType<FlatSymbolRefAttr>(
      TF::kStablehloEntryFunctionAttrName);
}

  }

  LogicalResult matchAndRewrite(TF::PartitionedCallOp op,
                                PatternRewriter& rewriter) const override {
    const auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(op.getFAttr());
    // Non-quantizable op
    if (!op->hasAttr(kQuantTraitAttrName)) return failure();
    StringRef function_name = f_attr.getValue();
    // TODO(b/228928859): Improve the getter function to match attributes rather