static bool AllowDynamicRangeQuantizedOperand(
      Operation* quantized_op, const CustomMap& custom_op_map) {
    auto call_op = cast<TF::PartitionedCallOp>(quantized_op);
    StringRef function_name =
        call_op.getFAttr().cast<FlatSymbolRefAttr>().getValue();
    // The below can be generalized as there are more read-only ops added such
    // as slice.
    const bool is_gather = function_name.contains("gather");

      devices.push_back(builder.getNamedAttr(name, metadata));
    }
  }

  op->setAttr(kDevicesAttr, builder.getDictionaryAttr(devices));
}

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

}

// 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);
}

              getElementTypeOrSelf(operand.getType()))) {
        legalize_to_tfl = false;
        return WalkResult::interrupt();
      }
    }
    return WalkResult::advance();
  });
  module->setAttr(kLegalizeTflVariables,
                  BoolAttr::get(context, legalize_to_tfl));
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateAnalyzeVariablesPass() {

// Set the (array of) exported name(s) of a (public) function to just
// contain the given string.
void SetExportedName(func::FuncOp f, StringRef name) {
  OpBuilder b(f);
  f->removeAttr(kTfSavedModelExportedNamesAttr);
  f->setAttr(kTfSavedModelExportedNamesAttr, b.getStrArrayAttr({name}));
}

// Convert a savedmodel public function into a private function.
// This means we need to remove any attributes that are only allowed

        input_specs.push_back(TypeAttr::get(requantize.params));
      } else {
        input_specs.push_back(TypeAttr::get(state.params));
      }
    }
    op->setAttr("input_specs", ArrayAttr::get(context, input_specs));

    llvm::SmallVector<Attribute, 4> output_specs;
    auto original_output_specs = op.getOutputSpecs().getValue();

                      /*arg_types=*/{}, /*arg_locs=*/{});

  init_func->setAttr(kTfSavedModelExportedNamesAttr,
                     builder.getStrArrayAttr(
                         {"tf_saved_model.session_initializer_restore_op"}));
  init_func->setAttr(
      kTfSavedModelInitializerTypeAttr,
      builder.getStringAttr(kTfSavedModelInitializerRestoreType));

      state = solver.lookupState<Device>(result);
      if (state) break;
    }
    if (!state || !state->hasDevice()) {
      return WalkResult::advance();
    }
    op->setAttr(kDevice, state->getDevice());
    return WalkResult::advance();
  });
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>
CreateColocateTPUCopyWithDynamicShapePass() {

                                PatternRewriter& rewriter) const override {
    auto input_op = op.getArg().getDefiningOp();
    if (auto q = llvm::dyn_cast_or_null<quantfork::QuantizeCastOp>(input_op)) {
      if (!q->getAttr(kVolatileOpAttrName)) return failure();

      if (remove_volatile_ops_type == kPreserveInputsAndOutputs) {
        // Don't remove leading and trailing QDQ for PTQ workflow, so the io
        // modifying lib can work correctly.

  auto func = builder.create<mlir::func::FuncOp>(
      ops[0]->getLoc(), dialect.str() + std::to_string(function_id),
      builder.getFunctionType(input_types, output_types));
  func->setAttr("dialect", builder.getStringAttr(dialect));
  auto block = func.addEntryBlock();

  llvm::DenseSet<Operation*> all_operations;
  for (Operation* outer : ops) {