add prod dtype

oneflow/core/functional/functional_api.yaml

@@ -274,7 +274,7 @@
   bind_python: True
 
 - name: "reduce_prod"
-  signature: "Tensor (Tensor x, Int32List axis, Bool keepdims=False) => ReduceProd"
+  signature: "Tensor (Tensor x, Int32List axis, Bool keepdims=False, *, DataType dtype=None) => ReduceProd"
   bind_python: True
 
 - name: "reduce_min_device_stage"

oneflow/core/functional/impl/math_functor.cpp

@@ -658,27 +658,29 @@ class ReduceProdFunctor {
         one::OpBuilder("reduce_prod").Input("input_tensor").Output("output_tensor").Build());
   }
   Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const std::vector<int32_t>& axis,
-                           const bool& keepdims) const {
+                           const bool& keepdims, const Optional<Symbol<DType>>& dtype) const {
     MutableAttrMap attrs;
+    std::shared_ptr<one::Tensor> tensor = x;
+    if (dtype.has_value() && (dtype != x->dtype())) { tensor = JUST(Cast(tensor, JUST(dtype))); }
     TensorProcessor tensor_processor;
     Symbol<DType> lowest_dtype;
-    if (DType::priority_order[x->dtype()->data_type()]
+    if (DType::priority_order[tensor->dtype()->data_type()]
         == DType::priority_order[DType::Bool()->data_type()]) {
       lowest_dtype = DType::Int64();
     } else {
-      lowest_dtype = x->dtype();
+      lowest_dtype = tensor->dtype();
     }
-    JUST(tensor_processor.AddInputs({x}, lowest_dtype).Apply());
+    JUST(tensor_processor.AddInputs({tensor}, lowest_dtype).Apply());
     TensorTuple input_tuple = JUST(tensor_processor.GetInputs());
     if (axis.empty()) {
-      std::vector<int32_t> reduce_axis(x->shape()->NumAxes());
+      std::vector<int32_t> reduce_axis(tensor->shape()->NumAxes());
       std::iota(reduce_axis.begin(), reduce_axis.end(), 0);
       JUST(attrs.SetAttr<std::vector<int32_t>>("axis", reduce_axis));
     } else {
       JUST(attrs.SetAttr<std::vector<int32_t>>("axis", axis));
     }
     JUST(attrs.SetAttr<bool>("keepdims", keepdims));
-    return OpInterpUtil::Dispatch<Tensor>(*op_, input_tuple, attrs);
+    return JUST(OpInterpUtil::Dispatch<Tensor>(*op_, input_tuple, attrs));
   }
 
  private:

python/oneflow/nn/modules/reduce_ops.py

@@ -36,11 +36,11 @@ def mean_op(input, dim=None, keepdim=False):
     return flow._C.reduce_mean(input, axis=axis_checked, keepdims=keepdim)
 
 
-def prod_op(input, dim=None, keepdim=False):
+def prod_op(input, dim=None, keepdim=False, *, dtype=None):
     axis_checked = _check_axis(dim, input.shape)
     if len(axis_checked) == 0:
         return input
-    return flow._C.reduce_prod(input, axis_checked, keepdim)
+    return flow._C.reduce_prod(input, axis_checked, keepdim, dtype=dtype)
 
 
 def all_op(input, dim=None, keepdim=False):

python/oneflow/test/modules/test_prod.py

@@ -51,6 +51,16 @@ def test_reduce_prod_bool_without_dim(test_case):
 
         return y
 
+    @autotest(auto_backward=False, check_graph=False)
+    def test_reduce_prod_with_dtype(test_case):
+        device = random_device()
+        ndim = random(1, 5).to(int)
+        x = random_tensor(ndim=ndim, low=1.0, high=4.0, requires_grad=False).to(device)
+        dim = random(0, ndim).to(int)
+        y = torch.prod(x, dim, dtype=torch.int32)
+
+        return y
+
 
 if __name__ == "__main__":
     unittest.main()