Missing Shape Inference for Prod (#4003)

Added required but missing Shape inference for `aten.prod`
llvm · Feb 7, 2025 · 52299e6 · 52299e6
1 parent d40d74f
commit 52299e6
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/projects/ltc/csrc/base_lazy_backend/shape_inference.cpp b/projects/ltc/csrc/base_lazy_backend/shape_inference.cpp
@@ -252,6 +252,18 @@ std::vector<torch::lazy::Shape> compute_shape_native_group_norm(
   return shapes;
 }
 
+std::vector<torch::lazy::Shape>
+compute_shape_prod(const at::Tensor &self,
+                   c10::optional<at::ScalarType> dtype) {
+  if (dtype.has_value()) {
+    return {Shape(dtype.value(), {})};
+  }
+  if (isIntegralType(self.scalar_type(), true)) {
+    return {Shape(c10::ScalarType::Long, {})};
+  }
+  return {Shape(self.scalar_type(), {})};
+}
+
 std::vector<torch::lazy::Shape>
 compute_shape_im2col(const at::Tensor &self, at::IntArrayRef kernel_size,
                      at::IntArrayRef dilation, at::IntArrayRef padding,