diff --git a/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td
index 99a6d08..92ae4d2 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.td
@@ -79,7 +79,8 @@ class HLOClient_BroadcastBinaryElementwiseOp<
   string mnemonic, list<OpTrait> traits> :
         HLOClient_Op<mnemonic,
             !listconcat(traits, [
-              DeclareOpInterfaceMethods<InferShapedTypeOpInterface>])> {
+              DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+              ["reifyReturnTypeShapes"]>])> {
   let arguments = (ins
     HLO_Tensor:$lhs,
     HLO_Tensor:$rhs,
@@ -99,13 +100,6 @@ class HLOClient_BroadcastBinaryElementwiseOp<
     $lhs `,` $rhs attr-dict `:`
     `(` type($lhs) `,` type($rhs) `)` `->` type(results)
   }];
-
-  let extraClassDeclaration = [{
-    // TODO(laurenzo): It isn't clear to me why reifyReturnShapes does not
-    // have its declaration generated by DeclareOpInterfaceMethods.
-    LogicalResult reifyReturnTypeShapes(
-         OpBuilder& builder, SmallVectorImpl<Value>& reifiedReturnShapes);
-  }];
 }
 
 def HLOClient_BroadcastAddOp : HLOClient_BroadcastBinaryElementwiseOp<"broadcast_add",