static_cast<float>(max.getValueAsDouble())});
    ElementsAttr axis_stats;
    IntegerAttr axis;

    quantfork::StatisticsOp stats_op = rewriter.create<quantfork::StatisticsOp>(
        op->getLoc(), op.getInput(), layer_stats, axis_stats, axis);
    op.getOutput().replaceAllUsesWith(stats_op.getResult());
    return success();
  }
};

static PassRegistration<ConvertCustomAggregationOpToQuantStatsPass> pass;

  const auto composite_resource_users =
      GetResourceOpsUsingCompositeArgsInReplicate(replicate_op);
  for (auto resource_user : composite_resource_users) {
    WrapOpInLaunch(builder, resource_user->getLoc(), resource_user,
                   tensorflow::GetDeviceAliasForLogicalCore(0));
  }
}

void TPUColocateCompositeResourceOps::runOnOperation() {

  if (call_op_has_platform_index_arg) {
    arg_mapper.map(func_op.getArgument(0),
                   builder.create<mhlo::ConstantOp>(
                       func_op.getLoc(), builder.getI16IntegerAttr(0)));
  }
  for (auto [func_arg, operand] : llvm::zip_equal(
           func_op.getArguments().take_back(call_op.getNumOperands()),
           call_op.getOperands())) {

  if (result_tuple_ty.size() != 1) return;
  auto result_ty = mlir::dyn_cast<ShapedType>(result_tuple_ty.getType(0));

  func.eraseBody();
  func.addEntryBlock();
  ImplicitLocOpBuilder builder(func.getLoc(), func.getBody());
  llvm::SmallVector<int32_t> result_shape_i32;
  auto result_shape = result_ty.getShape();
  for (auto element : result_shape) {
    result_shape_i32.push_back(static_cast<int32_t>(element));
  }

      return QuantizableOpsInDefaultMode().count(op_name) > 0;
    } else if (mode_ == kLegacyIntegerMode) {
      return QuantizableOpsInLegacyMode().count(op_name) > 0;
    } else {
      mlir::emitError(getOperation().getLoc(), "Unregconized mode: " + mode_);
      signalPassFailure();
      return true;
    }
  }

  // Converts the operation to a TFLite Flex op.
  bool ConvertToFlexOp(Operation *op);
};

    if (batch_norm.getIsTraining()) return failure();
    if (!batch_norm.getY().hasOneUse()) return failure();

    // Build the newly fused operation to replace the batch norm
    OperationState state(batch_norm.getLoc(),
                         _FusedBatchNormExOp::getOperationName());
    state.addOperands(batch_norm.getOperands());
    if (side_input) state.operands.push_back(side_input);
    state.addTypes(batch_norm.getResultTypes());

      TF::PartitionedCallOp call_op) const {
    // If unit-wise quantization config is found, overwrite the default config.
    auto quantization_unit = FindQuantizationUnitFromLoc(call_op.getLoc());
    if (!quantization_unit.has_value()) return std::nullopt;

    for (const auto& unit_config :
         quant_options_.unit_wise_quantization_specs()) {
      for (const auto& unit : unit_config.unit()) {

    std::replace(func_name.begin(), func_name.end(), '/', '_');

    FunctionType func_type =
        FunctionType::get(context, input_types, result_types);
    Location loc = metadata.ops.front()->getLoc();
    func::FuncOp func_op = func::FuncOp::create(loc, func_name, func_type);
    // Sets the device attribute for every input and every result of the
    // function.

def TrueBoolAttr : AttrConstraint<CPred<"$_self.cast<BoolAttr>().getValue()">>;

def CastValueToI64: NativeCodeCall<
  "CastValueToI64($0.getLoc(), $1, &$_builder)">;

def CastValueToElementType: NativeCodeCall<
  "$_builder.create<ConvertOp>($0.getLoc(), $1, "
  "getElementTypeOrSelf($2.getType()))">;

// Here, $0 is an ElementsAttr with exactly one element of type integer. $1 is

    const DictionaryAttr main_arg_attr =
        src_func_op.getArgAttrDict(src_arg_idx);

    main_func_op.insertArgument(main_arg_idx, src_arg.getType(), main_arg_attr,
                                src_arg.getLoc());

    const std::string new_input_name =
        absl::StrCat(GetInitializerType(src_func_op), "_", src_arg_idx, ":0");

    AddEntryFunctionInput(new_input_name, main_func_op);