diff --git a/paddle/framework/ddim.cc b/paddle/framework/ddim.cc index 3f949a6595ea326b97ac567daf9b35a68c8cf7f8..a1ae079f4abc21bef2614b4dd12e1713fb481447 100644 --- a/paddle/framework/ddim.cc +++ b/paddle/framework/ddim.cc @@ -178,13 +178,16 @@ std::vector vectorize(const DDim& ddim) { return result; } -ssize_t product(const DDim& ddim) { - ssize_t result = 1; - std::vector v = vectorize(ddim); - for (auto i : v) { - result *= i; +struct ProductVisitor : public boost::static_visitor { + template + ssize_t operator()(const Dim& dim) { + return product(dim); } - return result; +}; + +ssize_t product(const DDim& ddim) { + ProductVisitor visitor; + return boost::apply_visitor(visitor, ddim); } ///\cond HIDDEN diff --git a/paddle/framework/ddim_test.cc b/paddle/framework/ddim_test.cc index 36eef02370e0196c2af2c05f49176b70ce69235a..8ce7886f8ad952a0faf6413eb164b95a9ef1c680 100644 --- a/paddle/framework/ddim_test.cc +++ b/paddle/framework/ddim_test.cc @@ -52,6 +52,9 @@ TEST(DDim, Equality) { // product of a DDim EXPECT_EQ(paddle::framework::product(vddim), 45); + EXPECT_EQ( + paddle::framework::product(paddle::framework::make_ddim({3, 2, 5, 3})), + 90); } TEST(DDim, Print) {