block_arg.getLoc(),
      DenseIntElementsAttr::get(
          RankedTensorType::get({1, num_replicas}, builder.getIntegerType(32)),
          group_assignment_val));

  StringAttr reduce_op = builder.getStringAttr("Add");
  StringAttr mode = builder.getStringAttr("CrossReplica");
  return builder.create<XlaAllReduceOp>(block_arg.getLoc(), block_arg.getType(),

    auto *new_ewise_op =
        rewriter.create(op->getLoc(), op->getName().getIdentifier(),
                        {tpose_arg1.getOperand(0), tpose_arg2.getOperand(0)},
                        new_out_type, op->getAttrs());

    // Apply original tranpose to output of ewise op.
    auto out_tpose_op = rewriter.create<TFL::TransposeOp>(
        new_ewise_op->getLoc(), op->getResult(0).getType(),

  OpBuilder builder(func_.getBody());
  auto op = builder.create<mlir::TFL::NonMaxSuppressionV4Op>(
      func_.getLoc(), output_type0, output_type1, boxes, scores,
      max_output_size, iou_threshold, score_threshold);

  builder.create<mlir::func::ReturnOp>(func_.getLoc(), op.getResults());
}

LogicalResult ConvertNMSPaddedFunc::VerifySignature() {
  // Verify high-level function signature.

            const_op->getLoc(),
            DenseElementsAttr::get(
                tensorflow::GetTypeFromTFTensorShape(
                    {}, splat_elements_attr.getType().getElementType()),
                splat_elements_attr.getSplatValue<Attribute>()));
    TFL::FillOp fill = op_builder->create<TFL::FillOp>(
        const_op->getLoc(), splat_elements_attr.getType(), fill_shape,
        fill_value);

        CreateI8F32UniformQuantizedType(
            uniform_quantize_call_op.getLoc(), *rewriter.getContext(),
            input_scale_value, input_zero_point_value);
    auto input_uniform_quantize_op =
        rewriter.create<stablehlo::UniformQuantizeOp>(
            uniform_quantize_call_op.getLoc(),
            /*result=*/
            mlir::cast<TensorType>(input_value.getType())

    } else {
      initial_value = mlir::Attribute();
    }
    opToName[globalTensor] = name;
    auto variableOp = globalBuilder.create<ml_program::GlobalOp>(
        globalTensor.getLoc(), name, globalTensor.getType(),
        globalTensor.getIsMutable(), initial_value,
        /*visibility=*/globalBuilder.getStringAttr("private"));
    variableOp.setPrivate();
  }

  SymbolTable syms(module);

  SendTPUEmbeddingGradientsOp send_op;
  if (failed(GetOp(region, &send_op))) return failure();

  // No TPU embedding ops.
  if (!recv_op && !send_op) return success();

  Location loc = recv_op ? recv_op.getLoc() : send_op.getLoc();
  StringRef config = recv_op ? recv_op.getConfig() : send_op.getConfig();

  // Create XlaRecvTPUEmbeddingDeduplicationData op.
  OpBuilder builder(region);
  auto output_ty =

  void runOnOperation() override;
};

void ReplaceClusterReturnWithReturn(tf_device::ReturnOp cluster_return_op,
                                    OpBuilder* builder) {
  builder->create<func::ReturnOp>(cluster_return_op.getLoc(),
                                  cluster_return_op.getOperands());
  cluster_return_op.erase();
}

// Builds a function that outlines region attached to cluster_op or launch_op,

  SmallVector<Value, 4> tuple_operands;
  for (auto idx : llvm::seq<int32_t>(0, end))
    tuple_operands.push_back(
        builder.create<GetTupleElementOp>(value.getLoc(), value, idx)
            .getResult());

  return CreateTuple(builder, value.getLoc(), tuple_operands);
}

// Replaces uses of `values` with `replacements`. If `flatten_tuple` is false,

    // TODO: Map to a qbarrier with an attribute like [Forced] to signal that
    // this is a forced/hard-coded constraint.
    auto qbarrier = rewriter.create<QuantizeCastOp>(op.getLoc(), quantizedType,
                                                    op.getInputs());
    rewriter.replaceOpWithNewOp<DequantizeCastOp>(op, converter.input_type,