device_cluster.getLoc(), num_regions, llvm::ArrayRef<Type>{});
  SmallVector<Operation*, 4> core_to_host_insertion_point;
  SmallVector<mlir::tf_device::LaunchOp, 4> core_to_tmp_launch;
  SmallVector<Operation*, 4> compilation_key_ops;
  SmallVector<Value, 4> core_to_compilation_key;
  SmallVector<Operation*, 4> core_to_device_ordinal_op;
  SmallVector<Value, 4> core_to_device_ordinal;

      return failure();
    }

    SmallVector<int32_t, 4> begin, end, strides;
    SmallVector<int32_t, 4> padded_begin, padded_end, padded_strides;

    int num_input_dims = ranked_input_type.getRank();
    SmallVector<int32_t, 4> padding_begin(num_input_dims, 0);
    auto input_shape = ranked_input_type.getShape();
    SmallVector<int32_t, 4> padding_end(input_shape.begin(), input_shape.end());

    }
    auto body_attr =
        mlir::SymbolRefAttr::get(builder.getContext(), func_names.at(body_idx));

    return llvm::SmallVector<mlir::NamedAttribute, 4>{
        builder.getNamedAttr("cond", cond_attr),
        builder.getNamedAttr("body", body_attr)};
  }
  return llvm::SmallVector<mlir::NamedAttribute, 4>{};
}

Status ConvertSubgraphIdxToStablehloRegion(

    SmallVector<Type, 4> block_shape_splits_types(
        block_rank,
        tensorflow::GetTypeFromTFTensorShape({1}, rewriter.getIntegerType(64)));
    SmallVector<Value, 4> block_shape_splits(
        rewriter
            .create<SplitOp>(loc, block_shape_splits_types, zero_i32,
                             block_shape_i64)
            .getOutput());

    SmallVector<int64_t, 4> outer_shape_ints;

    llvm::SmallSet<int, 4> resized_tensor_lists =
        GetResizedTensorListIndexes(op.else_function(), tensor_list_args);

    llvm::SmallVector<Type, 8> result_types;
    result_types.reserve(op.getNumResults());
    llvm::SmallVector<Type, 4> op_result_types;
    for (Type ty : op.getResultTypes()) {
      op_result_types.push_back(ty);
    }

  int num_old_results = op.getNumResults();
  int num_new_results = hoister.GetLiftedNumResults();
  int num_extra_results = num_new_results - num_old_results;

  SmallVector<Type, 4> new_result_types;
  SmallVector<Value, 4> new_while_operands;
  new_result_types.resize(num_extra_results);
  new_while_operands.resize(num_extra_results);

  for (auto& it : hoister.GetResources()) {

  TensorType scale_type = RankedTensorType::get(
      {static_cast<int64_t>(num_channels)}, rewriter.getF32Type());
  TensorType zero_point_type = scale_type.clone(rewriter.getI32Type());

  llvm::SmallVector<float, 4> float_scales;
  llvm::SmallVector<int32_t, 4> int32_zero_points;
  float_scales.reserve(num_channels);
  int32_zero_points.reserve(num_channels);
  for (int i = 0; i < num_channels; ++i) {
    float_scales.push_back(scales[i]);

        combined_scale_multiply_op.getOperand(1).getDefiningOp());
    auto combined_scale_constant_op = cast<stablehlo::ConstantOp>(
        scale_combined_broadcast_in_dim_op.getOperand().getDefiningOp());

    SmallVector<double> filter_scale_values;
    for (const auto combined_scale_value :
         mlir::cast<DenseFPElementsAttr>(combined_scale_constant_op.getValue())
             .getValues<float>()) {

  // generalizing to multiple TPU clusters.
  Region* region = (*forward_pass_ops.begin())->getParentRegion();
  StringAttr replication_attr = GetReplicationAttr(*forward_pass_ops.begin());
  llvm::SmallVector<Operation*> checkset(forward_pass_ops.getArrayRef());
  checkset.append(backward_pass_ops.begin(), backward_pass_ops.end());
  for (Operation* op : checkset) {
    if (op->getParentRegion() != region) {