Update

ZhiweiYan-96 · ZhiweiYan-96 · commit e55b68ddee83 · 2024-11-11T08:46:11.000Z
[ghstack-poisoned]
diff --git a/aten/src/ATen/native/mkldnn/xpu/Blas.cpp b/aten/src/ATen/native/mkldnn/xpu/Blas.cpp
@@ -50,6 +50,11 @@ Tensor& addmm_out(
       mat1.dtype(),
       " != ",
       mat2.dtype())
+  // complex/double case
+  if (mat1.is_complex() || mat1.scalar_type() == ScalarType::Double) {
+    TORCH_CHECK(
+        false, "Double and complex datatype matmul is not supported in oneDNN");
+  }
 
   std::vector<int64_t> result_shape = {mat1.size(0), mat2.size(1)};
   result.resize_(result_shape);
@@ -77,12 +82,6 @@ Tensor& addmm_out(
       " but got:",
       self.sizes());
 
-  // complex/double case
-  if (mat1.is_complex() || mat1.scalar_type() == ScalarType::Double) {
-    TORCH_CHECK(
-        false, "Double and complex datatype matmul is not supported in oneDNN");
-  }
-
   // general case
   Tensor bias = Tensor();
   onednn::Attr attr;
@@ -292,6 +291,10 @@ Tensor& addbmm_out(
       batch1.dim(),
       " and ",
       batch2.dim());
+  if (self.is_complex() || self.scalar_type() == ScalarType::Double) {
+    TORCH_CHECK(
+        false, "Double and complex datatype matmul is not supported in oneDNN");
+  }
 
   out.resize_({batch1.size(1), batch2.size(2)});
   if (alpha.to<float>() == 0.f || batch1.numel() == 0 || batch2.numel() == 0) {
diff --git a/test/inductor/test_torchinductor_opinfo.py b/test/inductor/test_torchinductor_opinfo.py
@@ -343,7 +343,7 @@ def format_op(op):
     "cholesky_inverse": {f64},
     # could not create a primitive
     "addbmm": {f64},
-    "addmm": {f32, f64},
+    "addmm": {f64},
     "addmv": {f64},
     # could not create a primitive descriptor for
     # a deconvolution forward propagation primitive
