uint32_t num_replicas = replicate.getN();
  llvm::SmallVector<int32_t, 4> group_assignment_val;
  for (int i = 0; i < num_replicas; ++i) group_assignment_val.push_back(i);
  Value group_assignment = builder.create<ConstOp>(
      block_arg.getLoc(),
      DenseIntElementsAttr::get(
          RankedTensorType::get({1, num_replicas}, builder.getIntegerType(32)),
          group_assignment_val));

  // then update the OperandSegmentSize attribute.
  if (send_op) {
    int32_t operand_sizes[] = {static_cast<int32_t>(send_op.getN()),
                               static_cast<int32_t>(send_op.getNN()), 1};
    auto operand_size_attr = builder.getDenseI32ArrayAttr(operand_sizes);
    NamedAttrList attrs(send_op->getAttrs());

  for (llvm::StringRef device : device_list) {
    state_vars.push_back(builder->create<TF::VarHandleOp>(
        loc,
        llvm::ArrayRef<Type>{RankedTensorType::get(
            {}, TF::ResourceType::get(llvm::ArrayRef<TensorType>{key_type},
                                      builder->getContext()))},
        llvm::ArrayRef<Value>{},
        llvm::ArrayRef<NamedAttribute>{

// invariant. Shape ops are rewritten to be invariant when possible, prior to
// hoisting ops.
void HoistReplicateInvariantOps(tf_device::ReplicateOp replicate_op) {
  const int num_replicas = replicate_op.getN();
  Block* replicate_block = &replicate_op.GetBody();

  // A `ShapeOp` that directly depends on a `tf_device.replicate` param and does
  // not have a virtual device is assumed to return the same shape across all

    auto& arg = packed_args.emplace_back();
    arg.ssaName = std::get<0>(argAndType);
    arg.type = std::get<1>(argAndType);
  }
  if (parser.parseRegion(body, packed_args)) return failure();

  // Add derived `operandSegmentSizes` attribute based on parsed operands.
  if (!result.attributes.get(kOperandSegmentSizesAttr)) {
    int32_t num_replicated_inputs = replicated_inputs.size() * n;

    operand.get().replaceUsesWithIf(
        reduced_launch_op.getResult(operand.getOperandNumber()),
        operand_not_in_launch);
  }

  // Handle pass through block arguments.
  for (OpOperand& operand :
       original_launch_op.GetBody().getTerminator()->getOpOperands()) {
    if (mlir::isa<BlockArgument>(operand.get())) {

      ModuleOp::create(mlir::UnknownLoc::get(entry_func.getContext()),
                       absl::StrCat("module_", module_name));
  auto parent_module = entry_func->getParentOfType<ModuleOp>();
  auto versions_attr = parent_module->getAttr(kVersionsAttr);
  if (!versions_attr)
    return parent_module.emitError(CreateMissingAttributeMsg(kVersionsAttr));

  module_for_func.get().getOperation()->setAttr(kVersionsAttr, versions_attr);

    // types are static.
    SmallVector<DenseIntElementsAttr, 4> indices;
    indices.reserve(op.getN());
    for (auto it : llvm::zip(op.getIndices(), op.getData())) {
      Value index = std::get<0>(it);
      Value data = std::get<1>(it);

      DenseIntElementsAttr index_attr;
      if (!matchPattern(index, m_Constant(&index_attr))) return failure();

  if (!HasRankAtLeast(op.getX(), 2)) {
    return op.emitOpError("requires lhs operand to have rank at least two");
  }
  if (!HasRankAtLeast(op.getY(), 2)) {
    return op.emitOpError("requires rhs operand to have rank at least two");
  }

  RankedTensorType x_ty = GetRankedTensorTypeForOperand(op.getX());
  RankedTensorType y_ty = GetRankedTensorTypeForOperand(op.getY());

  if (!x_ty || !y_ty) return success();

    // `stablehlo.recv`.
    effects.emplace_back(MemoryEffects::Write::get(),
                         ResourceEffects::Send::get());
    effects.emplace_back(MemoryEffects::Write::get(),
                         ResourceEffects::Recv::get());
    effects.emplace_back(MemoryEffects::Write::get(),
                         ResourceEffects::XlaHostCompute::get());
  }
}

LogicalResult XlaCallModuleOp::verifySymbolUses(