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();
}

        set_segment_sizes_attr(input_ranges, op_def->input_arg(),
                               mlir::OpTrait::AttrSizedOperandSegments<
                                   void>::getOperandSegmentSizeAttr());
      }

      if (tf_op->hasTrait<mlir::OpTrait::AttrSizedResultSegments>()) {
        // Add derived "result_segment_sizes" attr to the created operation.

  }

  if (inst->hasTrait<mlir::OpTrait::AttrSizedOperandSegments>()) {
    // TODO(b/146937733): Don't use <void> here.
    llvm::StringRef attr_name = mlir::OpTrait::AttrSizedOperandSegments<
        void>::getOperandSegmentSizeAttr();
    attrs_to_ignore.insert(attr_name.data());
  }

  if (inst->hasTrait<mlir::OpTrait::AttrSizedResultSegments>()) {
    // TODO(b/146937733): Don't use <void> here.

      auto sizes = mlir::DenseI32ArrayAttr::get(context, {1, 1, 1, 0});
      auto attr_name =
          StringAttr::get(context, mlir::OpTrait::AttrSizedOperandSegments<
                                       void>::getOperandSegmentSizeAttr());
      attrs.push_back(NamedAttribute(attr_name, sizes));
    }

    // Insert fused operation right before the BiasAdd operation to guarantee

           << func_attr.getValue();
  }

  return success();
}

void BatchFunctionOp::eraseArguments(const BitVector& erase_indices) {
  const StringRef operand_segment_size_attr = getOperandSegmentSizeAttr();
  auto operandSegmentSizes = getOperation()->getAttrOfType<DenseI32ArrayAttr>(
      operand_segment_size_attr);

  // `operandSegmentSizes` attribute indicates the sizes of the two

      set_segment_sizes_attr(input_ranges, node.op_def().input_arg(),
                             mlir::OpTrait::AttrSizedOperandSegments<
                                 void>::getOperandSegmentSizeAttr());
    }

    if (inner_op->hasTrait<mlir::OpTrait::AttrSizedResultSegments>()) {
      // Add derived "result_segment_sizes" attr to the created operation.

      // variadic arguments). Attributes like `batch_timeout_micros` will not
      // make it to the return array.
      auto segment_sizes = getOperation()->getAttrOfType<DenseI32ArrayAttr>(getOperandSegmentSizeAttr());
      return OperandRangeRange(getOperands(), segment_sizes).join();
    }
  }];
}

def TF_BatchMatMulOp : TF_Op<"BatchMatMul", [Pure, TF_SameOperandsAndResultElementTypeResolveRef]> {