Select (#8)

* Select * Fix code style * Fix select messages
iefode · Sep 16, 2020 · f0a5399 · f0a5399
1 parent 4eca370
commit f0a5399
Show file tree

Hide file tree

Showing 7 changed files with 397 additions and 386 deletions.
diff --git a/ngraph/core/reference/include/ngraph/runtime/reference/select.hpp b/ngraph/core/reference/include/ngraph/runtime/reference/select.hpp
@@ -32,11 +32,14 @@ namespace ngraph
                         const T* arg1,
                         const T* arg2,
                         T* out,
-                        size_t count) // TODO: using char for bool, is this right?
+                        size_t arg0_count,
+                        size_t arg1_count,
+                        size_t arg2_count,
+                        size_t arg3_count) // TODO: using char for bool, is this right?
             {
-                for (size_t i = 0; i < count; i++)
+                for (size_t i = 0; i < arg3_count; i++)
                 {
-                    out[i] = arg0[i] ? arg1[i] : arg2[i];
+                    out[i] = arg0[i % arg0_count] ? arg1[i % arg1_count] : arg2[i % arg2_count];
                 }
             }
 

diff --git a/ngraph/core/src/pass/constant_folding_select.cpp b/ngraph/core/src/pass/constant_folding_select.cpp
@@ -40,6 +40,9 @@ shared_ptr<op::Constant> fold_constant_select(const shared_ptr<op::Constant>& se
                                       t->get_data_ptr<T>(),
                                       f->get_data_ptr<T>(),
                                       data_ptr,
+                                      shape_size(selection->get_shape()),
+                                      shape_size(t->get_shape()),
+                                      shape_size(f->get_shape()),
                                       shape_size(out_shape));
     }
     else if (auto select_v1 = as_type_ptr<op::v1::Select>(select))