static_cast<int32_t>(send_op.getNN()), 1};
    auto operand_size_attr = builder.getDenseI32ArrayAttr(operand_sizes);
    NamedAttrList attrs(send_op->getAttrs());
    attrs.set(send_op.getOperandSegmentSizeAttr(), operand_size_attr);

    AddOperandAndRewriteAs<XlaSendTPUEmbeddingGradientsOp>(send_op, dedup_op,
                                                           attrs, &builder);
  }
  return success();
}