auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);
  int start = -1, stop = -1;
  auto status = cc_ctx->op_kernel().InputRange(name, &start, &stop);
  args->start = start;
  args->stop = stop;
  tensorflow::Set_TF_Status_from_Status(args->status, status);
}

TF_DataType TF_InputDatatype(TF_OpKernelContext* ctx, int index) {

  groups_to_only = 0;
  for (const auto& [group, branch] : from) {
    auto to_iter = to.find(group);
    if (to_iter == to.end()) {
      ++groups_from_only;
    } else {
      auto to_branch = to_iter->second;
      if (to_branch == branch) {
        ++groups_same_branch;
      } else {
        ++groups_different_branch;
      }
    }
  }
  for (const auto& [group, _] : to) {

  if (auto repinput = dyn_cast<TF::TPUReplicatedInputOp>(op)) {
    if (!ValidateReplicatedInput(repinput, num_replicas, tpu_replicate_attr))
      return false;
  }
  if (auto repoutput = dyn_cast<TF::TPUReplicatedOutputOp>(op)) {
    if (!ValidateReplicatedOutput(repoutput, num_replicas, tpu_replicate_attr))
      return false;
  }
  if (auto partinput = dyn_cast<TF::TPUPartitionedInputOp>(op)) {

  SmallVector<int64_t, 2> output_shape{1, 2};
  auto input_type = RankedTensorType::get(input_shape, builder->getF32Type());
  auto weight_type = RankedTensorType::get(weight_shape, builder->getF32Type());
  auto bias_type = RankedTensorType::get(bias_shape, builder->getF32Type());
  auto projection_type =
      RankedTensorType::get(projection_shape, builder->getF32Type());
  auto layer_norm_scale_type =

  for (const auto *def : defs) {
    auto op_name = def->getValueAsString("opName");
    int tail_optional_tensor = 0, tensor_number_max = 0;
    auto *arg_values = def->getValueAsDag("arguments");
    for (int i = 0, e = arg_values->getNumArgs(); i < e; ++i) {
      auto arg = arg_values->getArg(i);
      auto *arg_def = dyn_cast<DefInit>(arg);
      if (!arg_def) continue;

  llvm::SmallVector<mlir::TF::IdentityOp, 8> erase_list;

  for (auto [cluster_result_id, cluster_result] :
       llvm::enumerate(cluster.getResults())) {
    auto [replicated_output, _] =
        GetSingleUserOfType<mlir::TF::TPUReplicatedOutputOp>(cluster_result);
    if (replicated_output) {
      for (OpResult per_replica_result : replicated_output->getResults()) {
        auto [partitioned_output, id_ops] =

  while (!worklist.empty()) {
    auto func_op = worklist.pop_back_val();
    auto& func_res = per_function_results.find(func_op)->getSecond();
    // In-function propagation.
    if (failed(ComputeResourceDevicesInComputation(func_op, &func_res)))
      return signalPassFailure();

    // Propagation to callees.
    auto walk_res = func_op.walk([&](Operation* op) {
      if (auto while_op = dyn_cast<WhileOp>(op)) {

                        AttrValueMap* values) {
  AttrValue value;
  auto& type_list = *value.mutable_list();
  for (auto type : types) {
    DataType dtype;
    TF_RETURN_IF_ERROR(ConvertScalarTypeToDataType(type, &dtype));
    type_list.add_type(dtype);
  }

  auto result = values->insert({string(name), value});
  assert(result.second && "cannot have multiple attributes with the same name");

  ASSERT_TRUE(module_op);

  auto test_func =
      module_op->lookupSymbol<func::FuncOp>("same_scale_after_composite");
  ASSERT_THAT(test_func, NotNull());

  auto quantize_op = FindOperationOfType<quantfork::QuantizeCastOp>(test_func);
  EXPECT_FALSE(IsOpQuantizableStableHlo(quantize_op));

  auto dequantize_op =
      FindOperationOfType<quantfork::DequantizeCastOp>(test_func);

    for (auto [i, value] : llvm::enumerate(partitioned_output.getOutput())) {
      const auto& resource = packed_input ? inputs[0] : inputs[i];
      builder.create<TF::AssignVariableOp>(
          assign_var->getLoc(), /*resource=*/resource, /*value=*/value);
    }
    assign_var.erase();
  }

  for (OpOperand& operand : cluster_func->getOpOperands()) {