[ NonZero ] operator with constant folding and onnx_importer support (#4431)

* There is the files "non_zero.cpp & non_zero.hpp" for which operator non_zero is defined"
The two files were added: non_zero.cpp & non_zero.hpp

* The pull request present a study of extension of the NonZero operatos used for ML models:
1) The various type of files was modified: cpp & hpp, onnx files & ng files
2) Non_zero operator supports the onnx standard (there are added the files in which the operator is enabled: non_zero.cpp & non_zero.hpp)
i

* Operator description: refactor

* GPU emitter - is not implemented

* GPU emitter: v0 version

* NoN zero: the input and output is adjusted to output shape & input_shape

* Move the NonZero op to v3 namespace

* Move NonZero to opset3

* Correct shape inference for NonZero

* NonZero op constant folding for 0D and 1D inputs

* Constant folding for NonZero op

* Correct output shape for NonZero & scalars

* Helper function to test NonZero

* NonZero constant folding UT for floats

* Enable more data types in NonZero CF

* NonZero type prop tests

* NonZero constant folding tests (directly)

* Use is_constant instead of casting in UT

* NonZero op doxygen docs

* onnx_importer docs adjustment

* Correct version of the NonZero core op

* Disable NonZero in GPU backend

* Short circuit if all elems in data are identical

* find_indices() optimization

* Assert on the input shape in NonZeroElements

* CF of NonZero with all non-zero values

* NonZero CF test for scalars

* bool support in NonZero

* Missing include in NonZero CF

* Dont throw if NonZero CF fails

* Update src/ngraph/pass/constant_folding_non_zero.cpp
Co-Authored-By: Robert Kimball <robert.kimball@intel.com>

* Removing warning
Co-authored-by: tomdol <tomasz.dolbniak@intel.com>
Co-authored-by: Scott Cyphers <diyessi@users.noreply.github.com>
Co-authored-by: Robert Kimball <robert.kimball@intel.com>

src/ngraph/CMakeLists.txt

@@ -269,6 +269,8 @@ set (SRC
     op/negative.hpp
     op/non_max_suppression.cpp
     op/non_max_suppression.hpp
+    op/non_zero.cpp
+    op/non_zero.hpp
     op/not.cpp
     op/not.hpp
     op/not_equal.cpp
@@ -485,6 +487,7 @@ set (SRC
     pass/constant_folding_dyn_slice.cpp
     pass/constant_folding_gather.cpp
     pass/constant_folding_logical_reduction.cpp
+    pass/constant_folding_non_zero.cpp
     pass/constant_folding_one_hot.cpp
     pass/constant_folding_pad.cpp
     pass/constant_folding_quantize.cpp

src/ngraph/frontend/onnx_import/CMakeLists.txt

@@ -152,6 +152,8 @@ add_library(onnx_import STATIC
         op/not.hpp
         op/non_max_suppression.cpp
         op/non_max_suppression.hpp
+        op/non_zero.cpp
+        op/non_zero.hpp
         op/onehot.cpp
         op/onehot.hpp
         op/or.hpp

src/ngraph/frontend/onnx_import/op/non_zero.cpp

+//*****************************************************************************
+// Copyright 2017-2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#include <memory>
+
+#include "ngraph/opsets/opset3.hpp"
+#include "non_zero.hpp"
+
+namespace ngraph
+{
+    namespace onnx_import
+    {
+        namespace op
+        {
+            namespace set_1
+            {
+                NodeVector non_zero(const Node& node)
+                {
+                    const auto data = node.get_ng_inputs().at(0);
+                    return {std::make_shared<ngraph::opset3::NonZero>(data)};
+                }
+
+            } // namespace set_1
+
+        } // namespace op
+
+    } // namespace onnx_import
+
+} // namespace ngraph

src/ngraph/frontend/onnx_import/op/non_zero.hpp

+//*****************************************************************************
+// Copyright 2017-2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#pragma once
+
+#include "core/node.hpp"
+#include "ngraph/node.hpp"
+
+namespace ngraph
+{
+    namespace onnx_import
+    {
+        namespace op
+        {
+            namespace set_1
+            {
+                /// \brief Convert ONNX NonZero operation to an nGraph node.
+                ///
+                /// \param node   The ONNX node object representing this operation.
+                ///
+                /// \return The vector containing nGraph nodes producing output of ONNX NonZero
+                ///         operation.
+                NodeVector non_zero(const Node& node);
+
+            } // namespace set_1
+
+        } // namespace op
+
+    } // namespace onnx_import
+
+} // namespace ngraph

src/ngraph/frontend/onnx_import/ops_bridge.cpp

@@ -88,6 +88,7 @@
 #include "op/mul.hpp"
 #include "op/neg.hpp"
 #include "op/non_max_suppression.hpp"
+#include "op/non_zero.hpp"
 #include "op/not.hpp"
 #include "op/onehot.hpp"
 #include "op/or.hpp"
@@ -314,6 +315,7 @@ namespace ngraph
             REGISTER_OPERATOR("Mul", 7, mul);
             REGISTER_OPERATOR("Neg", 1, neg);
             REGISTER_OPERATOR("NonMaxSuppression", 1, non_max_suppression);
+            REGISTER_OPERATOR("NonZero", 1, non_zero);
             REGISTER_OPERATOR("Not", 1, logical_not);
             REGISTER_OPERATOR("Or", 1, logical_or);
             REGISTER_OPERATOR("OneHot", 1, onehot);

src/ngraph/op/non_zero.cpp

+//*****************************************************************************
+// Copyright 2017-2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#include "ngraph/op/non_zero.hpp"
+#include "ngraph/op/op.hpp"
+
+using namespace ngraph;
+using namespace std;
+
+constexpr NodeTypeInfo op::v3::NonZero::type_info;
+
+op::v3::NonZero::NonZero(const Output<Node>& arg)
+    : Op({arg})
+{
+    constructor_validate_and_infer_types();
+}
+
+bool ngraph::op::v3::NonZero::visit_attributes(AttributeVisitor& visitor)
+{
+    return true;
+}
+
+void op::v3::NonZero::validate_and_infer_types()
+{
+    const PartialShape& input_shape = get_input_partial_shape(0);
+    const auto input_et = get_input_element_type(0);
+
+    NODE_VALIDATION_CHECK(this,
+                          input_et.is_integral() || input_et.is_real(),
+                          "NonZero input data type needs to be a numeric type. Got: ",
+                          input_et);
+
+    set_output_type(0, element::i64, PartialShape{input_shape.rank(), Dimension::dynamic()});
+    set_input_is_relevant_to_shape(0);
+}
+
+shared_ptr<Node> op::v3::NonZero::clone_with_new_inputs(const OutputVector& new_args) const
+{
+    check_new_args_count(this, new_args);
+    return make_shared<v3::NonZero>(new_args.at(0));
+}

src/ngraph/op/non_zero.hpp

+//*****************************************************************************
+// Copyright 2017-2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#pragma once
+
+#include "ngraph/op/op.hpp"
+
+namespace ngraph
+{
+    namespace op
+    {
+        namespace v3
+        {
+            /// \brief NonZero operation returning indices of non-zero elements in the input tensor.
+            ///
+            /// \note The indices are returned by-dimension in row-major order. For example
+            ///       the following output contains 3 indices of a 3D input tensor elements:
+            ///       [[0, 0, 2],
+            ///        [0, 1, 1],
+            ///        [0, 1, 2]]
+            ///       The values point to input elements at [0,0,0], [0,1,1] and [2,1,2]
+            class NGRAPH_API NonZero : public Op
+            {
+            public:
+                static constexpr NodeTypeInfo type_info{"NonZero", 3};
+                const NodeTypeInfo& get_type_info() const override { return type_info; }
+                /// \brief Constructs a NonZero operation.
+                NonZero() = default;
+                /// \brief Constructs a NonZero operation.
+                ///
+                /// \param arg Node that produces the input tensor.
+                NonZero(const Output<Node>& arg);
+
+                bool visit_attributes(AttributeVisitor& visitor) override;
+                void validate_and_infer_types() override;
+
+                virtual std::shared_ptr<Node>
+                    clone_with_new_inputs(const OutputVector& new_args) const override;
+            };
+        }
+        using v3::NonZero;
+    } // namespace op
+} // namespace ngraph

src/ngraph/op/op_version_tbl.hpp

@@ -152,6 +152,7 @@ NGRAPH_OP(Multiply, ngraph::op::v0, 0)
 NGRAPH_OP(Multiply, ngraph::op::v1, 1)
 NGRAPH_OP(Negative, ngraph::op::v0, 0)
 NGRAPH_OP(NonMaxSuppression, ngraph::op::v1, 1)
+NGRAPH_OP(NonZero, ngraph::op::v3, 3)
 NGRAPH_OP(NormalizeL2, ngraph::op::v0, 0)
 NGRAPH_OP(Not, ngraph::op::v0, 0)
 NGRAPH_OP(NotEqual, ngraph::op::v0, 0)

src/ngraph/ops.hpp

@@ -135,6 +135,7 @@
 #include "ngraph/op/multiply.hpp"
 #include "ngraph/op/negative.hpp"
 #include "ngraph/op/non_max_suppression.hpp"
+#include "ngraph/op/non_zero.hpp"
 #include "ngraph/op/not.hpp"
 #include "ngraph/op/not_equal.hpp"
 #include "ngraph/op/one_hot.hpp"

src/ngraph/opsets/opset3.hpp

+//*****************************************************************************
+// Copyright 2017-2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#pragma once
+
+#include "ngraph/ops.hpp"
+
+namespace ngraph
+{
+    namespace opset3
+    {
+#define NGRAPH_OP(a, b) using b::a;
+#include "ngraph/opsets/opset3_tbl.hpp"
+#undef NGRAPH_OP
+    }
+}

src/ngraph/opsets/opset3_tbl.hpp

+//*****************************************************************************
+// Copyright 2017-2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#ifndef NGRAPH_OP
+#warning "NGRAPH_OP not defined"
+#define NGRAPH_OP(x, y)
+#endif
+
+#include "opset2_tbl.hpp"
+NGRAPH_OP(NonZero, ngraph::op::v3)

src/ngraph/pass/constant_folding.hpp

@@ -62,7 +62,8 @@ public:
         SPLIT,
         VARIADIC_SPLIT,
         ONE_HOT,
-        TILE
+        TILE,
+        NON_ZERO
     };
 
     ConstantFolding(const ngraph::BuildNodeExecutorMap& cfmap = ngraph::BuildNodeExecutorMap())
@@ -99,6 +100,7 @@ public:
         construct_constant_unsqueeze();
         construct_constant_one_hot();
         construct_constant_tile();
+        construct_constant_non_zero();
     }
 
     // this allows to specify the order in which matchers will be run
@@ -144,6 +146,7 @@ public:
             case CFTransformations::VARIADIC_SPLIT: construct_constant_variadic_split(); break;
             case CFTransformations::ONE_HOT: construct_constant_one_hot(); break;
             case CFTransformations::TILE: construct_constant_tile(); break;
+            case CFTransformations::NON_ZERO: construct_constant_non_zero(); break;
             }
         }
     }
@@ -177,6 +180,7 @@ private:
     void construct_constant_variadic_split();
     void construct_constant_one_hot();
     void construct_constant_tile();
+    void construct_constant_non_zero();
 
     ngraph::BuildNodeExecutorMap m_cfmap;
 };

src/ngraph/pass/constant_folding_non_zero.cpp

+//*****************************************************************************
+// Copyright 2017-2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#include <numeric>
+
+#include "constant_folding.hpp"
+#include "ngraph/op/non_zero.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+namespace
+{
+    struct NonZeroElements
+    {
+        using Results_t = std::vector<std::vector<int64_t>>;
+
+        NonZeroElements(const Shape& input_shape)
+            : m_input_shape{input_shape}
+            , m_results{Results_t(input_shape.size())}
+        {
+            NGRAPH_CHECK(m_input_shape.size() > 0,
+                         "Can't use the NonZeroElements class with a scalar shape");
+        }
+
+        template <typename T>
+        void find_indices(const T* values)
+        {
+            m_current_index = Shape(m_input_shape.size(), 0UL);
+            const auto values_count = shape_size(m_input_shape);
+
+            const T zero_value = T{0};
+            for (size_t i = 0; i + 1 < values_count; ++i)
+            {
+                if (values[i] != zero_value)
+                {
+                    add_to_results(m_current_index);
+                }
+
+                next_index();
+            }
+
+            // check the last element in the input values
+            if (values_count != 0 && values[values_count - 1] != zero_value)
+            {
+                add_to_results(m_current_index);
+            }
+        }
+
+        void generate_all_indices()
+        {
+            m_current_index = Shape(m_input_shape.size(), 0UL);
+            size_t i = 0;
+            const auto values_count = shape_size(m_input_shape);
+            while (i + 1 < values_count)
+            {
+                add_to_results(m_current_index);
+                next_index();
+                ++i;
+            }
+            add_to_results(m_current_index);
+        }
+
+        const Results_t& get_indices() const { return m_results; }
+    private:
+        /// \brief Adds single dimensions of an index into the matching element of the results
+        void add_to_results(const Shape& index)
+        {
+            for (size_t dim = 0; dim < index.size(); ++dim)
+            {
+                m_results.at(dim).push_back(index[dim]);
+            }
+        }
+
+        // Generates an index pointing to the next element in the flattened tensor
+        // It behaves similar to flipping bits when incrementing a binary number
+        inline void next_index()
+        {
+            for (size_t dim = m_current_index.size() - 1; dim >= 0; --dim)
+            {
+                auto& dim_value = m_current_index.at(dim);
+                if (dim_value + 1 == m_input_shape[dim])
+                {
+                    dim_value = 0;
+                }
+                else
+                {
+                    ++dim_value;
+                    return;
+                }
+            }
+        }
+
+    private:
+        const Shape m_input_shape;
+        Results_t m_results;
+        Shape m_current_index;
+    };
+}
+
+template <typename T>
+static shared_ptr<op::Constant> fold_constant_non_zero(const shared_ptr<op::Constant>& data)
+{
+    const auto input_shape = data->get_shape();
+    const auto* input_values = data->get_data_ptr<T>();
+    const bool identical_elems_in_data = data->get_all_data_elements_bitwise_identical();
+
+    if (identical_elems_in_data && input_values[0] == T{0})
+    {
+        return nullptr;
+    }
+
+    if (ngraph::is_scalar(input_shape))
+    {
+        return op::Constant::create(element::i64, Shape{1, 1}, {0});
+    }
+    else if (is_vector(input_shape))
+    {
+        const auto input_values_count = shape_size(input_shape);
+        std::vector<int64_t> indices;
+        indices.reserve(input_values_count);
+
+        if (identical_elems_in_data)
+        {
+            // return a complete set of indices since all of them are non-zero
+            indices.resize(input_values_count);
+            std::iota(indices.begin(), indices.end(), 0);
+        }
+        else
+        {
+            const T zero_value = T{0};
+            for (size_t i = 0; i < input_values_count; ++i)
+            {
+                if (input_values[i] != zero_value)
+                {
+                    indices.push_back(i);
+                }
+            }
+
+            indices.shrink_to_fit();
+        }
+
+        return op::Constant::create(element::i64, Shape{1, indices.size()}, indices);
+    }
+    else
+    {
+        NonZeroElements non_zero_elems{input_shape};
+
+        if (identical_elems_in_data)
+        {
+            non_zero_elems.generate_all_indices();
+        }
+        else
+        {
+            non_zero_elems.find_indices(input_values);
+        }
+
+        const auto& found_indices = non_zero_elems.get_indices();
+
+        // flatten the results and return them as a Constant
+        std::vector<int64_t> indices;
+        indices.reserve(found_indices.size() * found_indices.front().size());
+        for (const auto& row : found_indices)
+        {
+            indices.insert(indices.end(), row.begin(), row.end());
+        }
+
+        const Shape out_shape{found_indices.size(), found_indices.front().size()};
+        return op::Constant::create(element::i64, out_shape, indices);
+    }
+}
+
+void pass::ConstantFolding::construct_constant_non_zero()
+{
+    const auto data_label = make_shared<pattern::op::Label>(
+        element::f32, Shape{2, 2, 3}, pattern::has_class<op::Constant>());
+    const auto non_zero = make_shared<op::v3::NonZero>(data_label);
+
+    auto constant_non_zero_callback = [data_label](pattern::Matcher& m) {
+        auto pattern_map = m.get_pattern_map();
+
+        const auto data = static_pointer_cast<op::Constant>(pattern_map[data_label]);
+
+        std::shared_ptr<Node> replacement;
+        switch (data->get_element_type())
+        {
+        case element::Type_t::boolean: replacement = fold_constant_non_zero<char>(data); break;
+        case element::Type_t::bf16: replacement = fold_constant_non_zero<bfloat16>(data); break;
+        case element::Type_t::f16: replacement = fold_constant_non_zero<float16>(data); break;
+        case element::Type_t::f32: replacement = fold_constant_non_zero<float>(data); break;
+        case element::Type_t::f64: replacement = fold_constant_non_zero<double>(data); break;
+        case element::Type_t::i8: replacement = fold_constant_non_zero<int8_t>(data); break;
+        case element::Type_t::i16: replacement = fold_constant_non_zero<int16_t>(data); break;
+        case element::Type_t::i32: replacement = fold_constant_non_zero<int32_t>(data); break;
+        case element::Type_t::i64: replacement = fold_constant_non_zero<int64_t>(data); break;
+        case element::Type_t::u8: replacement = fold_constant_non_zero<uint8_t>(data); break;
+        case element::Type_t::u16: replacement = fold_constant_non_zero<uint16_t>(data); break;
+        case element::Type_t::u32: replacement = fold_constant_non_zero<uint32_t>(data); break;
+        case element::Type_t::u64: replacement = fold_constant_non_zero<uint64_t>(data); break;
+        case element::Type_t::u1:
+        case element::Type_t::dynamic:
+        case element::Type_t::undefined:
+            NGRAPH_CHECK(false, "Unsupported data type in NonZero constant folding");
+            break;
+        }
+
+        if (replacement.get() != nullptr)
+        {
+            replace_node(m.get_match_root(), replacement);
+            return true;
+        }
+        else
+        {
+            return false;
+        }
+    };
+
+    const auto matcher =
+        make_shared<pattern::Matcher>(non_zero, "ConstantFolding.ConstantNonZeroV3");
+    this->add_matcher(matcher, constant_non_zero_callback, PassProperty::CHANGE_DYNAMIC_STATE);
+}

src/ngraph/runtime/gpu/gpu_emitter.cpp

@@ -1990,3 +1990,8 @@ std::string runtime::gpu::GPU_Emitter::emit_v0_LayerNormBackprop(EMIT_ARGS)
 {
     throw unsupported_op("Unsupported op '" + node->description() + "'");
 }
+
+std::string runtime::gpu::GPU_Emitter::emit_v3_NonZero(EMIT_ARGS)
+{
+    throw unsupported_op("Unsupported op '" + node->description() + "'");
+}

src/ngraph/serializer.cpp

@@ -57,6 +57,7 @@ namespace
     {
 #define VSUF0(NAME) NAME
 #define VSUF1(NAME) NAME##_v1
+#define VSUF3(NAME) NAME##_v3
 #define NGRAPH_OP(NAME, NAMESPACE, VERSION) VSUF##VERSION(NAME),
 #include "ngraph/op/op_version_tbl.hpp"
 #undef NGRAPH_OP
@@ -2246,6 +2247,12 @@ shared_ptr<Node> JSONDeserializer::deserialize_node(json node_js)
 
             break;
         }
+        case OP_TYPEID::NonZero_v3:
+        {
+            node = make_shared<op::v3::NonZero>(args[0]);
+
+            break;
+        }
         case OP_TYPEID::NormalizeL2:
         {
             float eps = node_js.at("eps").get<float>();
@@ -4106,6 +4113,8 @@ json JSONSerializer::serialize_node(const Node& n)
         node["sort_result_descending"] = tmp->get_sort_result_descending();
         break;
     }
+    case OP_TYPEID::NonZero_v3: { break;
+    }
     case OP_TYPEID::NormalizeL2:
     {
         auto tmp = static_cast<const op::NormalizeL2*>(&n);

test/CMakeLists.txt

@@ -156,6 +156,7 @@ set(SRC
     type_prop/max_pool.cpp
     type_prop/mvn.cpp
     type_prop/non_max_suppression.cpp
+    type_prop/non_zero.cpp
     type_prop/normalize.cpp
     type_prop/one_hot.cpp
     type_prop/pad.cpp

test/constant_folding.cpp

@@ -2376,3 +2376,157 @@ TEST(constant_folding, pass_property)
     ASSERT_FALSE(pass->get_property(pass::PassProperty::REQUIRE_STATIC_SHAPE));
     ASSERT_TRUE(pass->get_property(pass::PassProperty::CHANGE_DYNAMIC_STATE));
 }
+
+TEST(constant_folding, constant_non_zero_0D)
+{
+    auto data = op::Constant::create(element::i32, Shape{}, {1});
+    auto non_zero = make_shared<op::v3::NonZero>(data);
+    auto f = make_shared<Function>(non_zero, ParameterVector{});
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(f);
+
+    ASSERT_EQ(count_ops_of_type<op::v3::NonZero>(f), 0);
+    ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
+
+    const auto new_const = as_type_ptr<op::Constant>(f->get_results().at(0)->get_argument(0));
+    ASSERT_TRUE(new_const);
+    const auto values_out = new_const->get_vector<int64_t>();
+
+    const vector<int64_t> values_expected{0};
+    ASSERT_EQ(values_expected, values_out);
+    ASSERT_EQ((Shape{1, 1}), new_const->get_shape());
+}
+
+TEST(constant_folding, constant_non_zero_1D)
+{
+    vector<int> values_in{0, 1, 0, 1};
+    auto data = make_shared<op::Constant>(element::i32, Shape{4}, values_in);
+    auto non_zero = make_shared<op::v3::NonZero>(data);
+    auto f = make_shared<Function>(non_zero, ParameterVector{});
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(f);
+
+    ASSERT_EQ(count_ops_of_type<op::v3::NonZero>(f), 0);
+    ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
+
+    const auto new_const = as_type_ptr<op::Constant>(f->get_results().at(0)->get_argument(0));
+    ASSERT_TRUE(new_const);
+    const auto values_out = new_const->get_vector<int64_t>();
+
+    const vector<int64_t> values_expected{1, 3};
+    ASSERT_EQ(values_expected, values_out);
+    ASSERT_EQ((Shape{1, 2}), new_const->get_shape());
+}
+
+TEST(constant_folding, constant_non_zero_1D_all_indices)
+{
+    const vector<float> values_in{1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f};
+    const auto data = make_shared<op::Constant>(element::f32, Shape{values_in.size()}, values_in);
+    const auto non_zero = make_shared<op::v3::NonZero>(data);
+    auto f = make_shared<Function>(non_zero, ParameterVector{});
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(f);
+
+    ASSERT_EQ(count_ops_of_type<op::v3::NonZero>(f), 0);
+    ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
+
+    const auto new_const = as_type_ptr<op::Constant>(f->get_results().at(0)->get_argument(0));
+    ASSERT_TRUE(new_const);
+    const auto values_out = new_const->get_vector<int64_t>();
+
+    const vector<int64_t> values_expected{0, 1, 2, 3, 4, 5, 6, 7};
+    ASSERT_EQ(values_expected, values_out);
+    ASSERT_EQ((Shape{1, values_in.size()}), new_const->get_shape());
+}
+
+TEST(constant_folding, constant_non_zero_2D)
+{
+    vector<int> values_in{1, 0, 0, 0, 1, 0, 1, 1, 0};
+    auto data = make_shared<op::Constant>(element::i32, Shape{3, 3}, values_in);
+    auto non_zero = make_shared<op::v3::NonZero>(data);
+    auto f = make_shared<Function>(non_zero, ParameterVector{});
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(f);
+
+    ASSERT_EQ(count_ops_of_type<op::v3::NonZero>(f), 0);
+    ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
+
+    const auto new_const = as_type_ptr<op::Constant>(f->get_results().at(0)->get_argument(0));
+    ASSERT_TRUE(new_const);
+    const auto values_out = new_const->get_vector<int64_t>();
+
+    const vector<int64_t> values_expected{0, 1, 2, 2, 0, 1, 0, 1};
+    ASSERT_EQ(values_expected, values_out);
+    ASSERT_EQ((Shape{2, 4}), new_const->get_shape());
+}
+
+TEST(constant_folding, constant_non_zero_2D_all_indices)
+{
+    const vector<int8_t> values_in{1, 1, 1, 1, 1, 1, 1, 1, 1};
+    const auto data = make_shared<op::Constant>(element::i8, Shape{3, 3}, values_in);
+    const auto non_zero = make_shared<op::v3::NonZero>(data);
+    auto f = make_shared<Function>(non_zero, ParameterVector{});
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(f);
+
+    ASSERT_EQ(count_ops_of_type<op::v3::NonZero>(f), 0);
+    ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
+
+    const auto new_const = as_type_ptr<op::Constant>(f->get_results().at(0)->get_argument(0));
+    ASSERT_TRUE(new_const);
+    const auto values_out = new_const->get_vector<int64_t>();
+
+    const vector<int64_t> values_expected{0, 0, 0, 1, 1, 1, 2, 2, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2};
+    ASSERT_EQ(values_expected, values_out);
+    ASSERT_EQ((Shape{2, values_in.size()}), new_const->get_shape());
+}
+
+TEST(constant_folding, constant_non_zero_2D_all_zeros)
+{
+    const vector<uint8_t> values_in{0, 0, 0, 0, 0, 0};
+    const auto data = make_shared<op::Constant>(element::u8, Shape{2, 3}, values_in);
+    const auto non_zero = make_shared<op::v3::NonZero>(data);
+    auto f = make_shared<Function>(non_zero, ParameterVector{});
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(f);
+
+    // constant folding should fail and the NonZero op should still be in the graph
+    ASSERT_EQ(count_ops_of_type<op::v3::NonZero>(f), 1);
+    ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
+}
+
+TEST(constant_folding, constant_non_zero_3D)
+{
+    vector<int> values_in{1, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 0};
+    auto data = make_shared<op::Constant>(element::i32, Shape{2, 3, 3}, values_in);
+    auto non_zero = make_shared<op::v3::NonZero>(data);
+    auto f = make_shared<Function>(non_zero, ParameterVector{});
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(f);
+
+    ASSERT_EQ(count_ops_of_type<op::v3::NonZero>(f), 0);
+    ASSERT_EQ(count_ops_of_type<op::Constant>(f), 1);
+
+    const auto new_const = as_type_ptr<op::Constant>(f->get_results().at(0)->get_argument(0));
+    ASSERT_TRUE(new_const);
+    const auto values_out = new_const->get_vector<int64_t>();
+
+    const vector<int64_t> values_expected{0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 2, 2, 2,
+                                          0, 0, 0, 1, 1, 2, 0, 2, 1, 0, 1, 2, 0, 1, 2, 0, 2, 1};
+    ASSERT_EQ(values_expected, values_out);
+    ASSERT_EQ((Shape{3, 12}), new_const->get_shape());
+}

test/models/onnx/non_zero_1d.prototxt

+ir_version: 3
+producer_name: "nGraph ONNX Importer"
+model_version: 1
+graph {
+  node {
+    name: "non_zero"
+    input: "A"
+    output: "out"
+    op_type: "NonZero"
+  }
+  input {
+    name: "A"
+    type {
+      tensor_type {
+        elem_type: 6
+        shape {
+          dim {
+            dim_value: 5
+          }
+        }
+      }
+    }
+  }
+  initializer {
+    data_type: 6
+    name: "A"
+    dims: 5
+    int32_data: 0
+    int32_data: 1
+    int32_data: 2
+    int32_data: 0
+    int32_data: 3
+  }
+  output {
+    name: "out"
+    type {
+      tensor_type {
+        elem_type: 7
+      }
+    }
+  }
+  name: "non_zero_1d"
+}
+opset_import {
+  domain: ""
+  version: 9
+}

test/models/onnx/non_zero_1d_float.prototxt

+ir_version: 3
+producer_name: "nGraph ONNX Importer"
+model_version: 1
+graph {
+  node {
+    name: "non_zero"
+    input: "A"
+    output: "out"
+    op_type: "NonZero"
+  }
+  input {
+    name: "A"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+            dim_value: 10
+          }
+        }
+      }
+    }
+  }
+  initializer {
+    data_type: 1
+    name: "A"
+    dims: 10
+    float_data: -1234.567
+    float_data: -0.5
+    float_data: 0.0
+    float_data: 0.1
+    float_data: 0.01
+    float_data: 0.001
+    float_data: 0.0001
+    float_data: 0.00001
+    float_data: 0.000001
+    float_data: 0.0000001
+  }
+  output {
+    name: "out"
+    type {
+      tensor_type {
+        elem_type: 7
+      }
+    }
+  }
+  name: "non_zero_1d_float"
+}
+opset_import {
+  domain: ""
+  version: 9
+}

test/models/onnx/non_zero_2d_bool.prototxt

+ir_version: 3
+producer_name: "nGraph ONNX Importer"
+model_version: 1
+graph {
+  node {
+    name: "non_zero"
+    input: "A"
+    output: "out"
+    op_type: "NonZero"
+  }
+  input {
+    name: "A"
+    type {
+      tensor_type {
+        elem_type: 9
+        shape {
+          dim {
+            dim_value: 2
+          }
+          dim {
+            dim_value: 2
+          }
+        }
+      }
+    }
+  }
+  initializer {
+    data_type: 9
+    name: "A"
+    dims: 2
+    dims: 2
+    int32_data: 0
+    int32_data: 1
+    int32_data: 1
+    int32_data: 0
+  }
+  output {
+    name: "out"
+    type {
+      tensor_type {
+        elem_type: 7
+      }
+    }
+  }
+  name: "non_zero_2d_bool"
+}
+opset_import {
+  domain: ""
+  version: 9
+}

test/models/onnx/non_zero_3d.prototxt

+ir_version: 3
+producer_name: "nGraph ONNX Importer"
+model_version: 1
+graph {
+  node {
+    name: "non_zero"
+    input: "A"
+    output: "out"
+    op_type: "NonZero"
+  }
+  input {
+    name: "A"
+    type {
+      tensor_type {
+        elem_type: 6
+        shape {
+          dim {
+            dim_value: 3
+          }
+          dim {
+            dim_value: 2
+          }
+          dim {
+            dim_value: 2
+          }
+        }
+      }
+    }
+  }
+  initializer {
+    data_type: 6
+    name: "A"
+    dims: 3
+    dims: 2
+    dims: 2
+    int32_data: 1
+    int32_data: 1
+    int32_data: 0
+    int32_data: 1
+    int32_data: 0
+    int32_data: 1
+    int32_data: 1
+    int32_data: 0
+    int32_data: 1
+    int32_data: 0
+    int32_data: 0
+    int32_data: 1
+  }
+  output {
+    name: "out"
+    type {
+      tensor_type {
+        elem_type: 7
+      }
+    }
+  }
+  name: "non_zero_3d"
+}
+opset_import {
+  domain: ""
+  version: 9
+}

test/models/onnx/non_zero_scalar.prototxt

+ir_version: 3
+producer_name: "nGraph ONNX Importer"
+model_version: 1
+graph {
+  node {
+    name: "non_zero"
+    input: "A"
+    output: "out"
+    op_type: "NonZero"
+  }
+  input {
+    name: "A"
+    type {
+      tensor_type {
+        elem_type: 6
+        shape {
+        }
+      }
+    }
+  }
+  initializer {
+    data_type: 6
+    name: "A"
+    int32_data: 777
+  }
+  output {
+    name: "out"
+    type {
+      tensor_type {
+        elem_type: 7
+      }
+    }
+  }
+  name: "non_zero_scalar"
+}
+opset_import {
+  domain: ""
+  version: 9
+}

test/onnx/onnx_import.in.cpp

@@ -38,6 +38,8 @@
 #include "ngraph/frontend/onnx_import/onnx.hpp"
 #include "ngraph/frontend/onnx_import/onnx_utils.hpp"
 #include "ngraph/ngraph.hpp"
+#include "ngraph/pass/manager.hpp"
+#include "ngraph/pass/constant_folding.hpp"
 #include "util/all_close.hpp"
 #include "util/all_close_f.hpp"
 #include "util/ndarray.hpp"
@@ -2052,3 +2054,88 @@ NGRAPH_TEST(onnx_${BACKEND_NAME}, model_round)
 
     test_case.run();
 }
+
+namespace
+{
+    void test_non_zero_constant_folding(std::shared_ptr<ngraph::Function> function,
+                                        const std::vector<int64_t> expected_output,
+                                        const PartialShape expected_shape = PartialShape::dynamic())
+    {
+        ngraph::pass::Manager pass_manager;
+        pass_manager.register_pass<pass::ConstantFolding>();
+        pass_manager.run_passes(function);
+
+        for (auto ng_node : function->get_ordered_ops())
+        {
+            if (ng_node->is_constant())
+            {
+                const auto folded_non_zero = as_type_ptr<op::Constant>(ng_node);
+                const auto values = folded_non_zero->cast_vector<int64_t>();
+
+                EXPECT_TRUE(ngraph::test::all_close(expected_output, values));
+
+                if (expected_shape.is_static())
+                {
+                    EXPECT_EQ(folded_non_zero->get_output_shape(0), expected_shape.to_shape());
+                }
+
+                return;
+            }
+        }
+
+        FAIL() << "NonZero constant folding failed.";
+    }
+}
+
+NGRAPH_TEST(onnx_${BACKEND_NAME}, model_non_zero_scalar)
+{
+    const auto fn = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/non_zero_scalar.prototxt"));
+
+    test_non_zero_constant_folding(fn, {0}, Shape{1, 1});
+}
+
+NGRAPH_TEST(onnx_${BACKEND_NAME}, model_non_zero_1d)
+{
+    const auto fn = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/non_zero_1d.prototxt"));
+
+    ngraph::pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(fn);
+
+    test_non_zero_constant_folding(fn, {1, 2, 4}, Shape{1, 3});
+}
+
+NGRAPH_TEST(onnx_${BACKEND_NAME}, model_non_zero_1d_float)
+{
+    const auto fn = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/non_zero_1d_float.prototxt"));
+
+    ngraph::pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ConstantFolding>();
+    pass_manager.run_passes(fn);
+
+    test_non_zero_constant_folding(fn, {0, 1, 3, 4, 5, 6, 7, 8, 9});
+}
+
+NGRAPH_TEST(onnx_${BACKEND_NAME}, model_non_zero_3d)
+{
+    const auto fn = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/non_zero_3d.prototxt"));
+
+    // Vertical slices are 3D indices of non-zero elements in the input tensor
+    // {0, 0, 0, 1, 1, 2, 2}
+    // {0, 0, 1, 0, 1, 0, 1}
+    // {0, 1, 1, 1, 0, 0, 1}
+    test_non_zero_constant_folding(fn,
+                                   {0, 0, 0, 1, 1, 2, 2, 0, 0, 1, 0, 1, 0, 1, 0, 1, 1, 1, 0, 0, 1});
+}
+
+NGRAPH_TEST(onnx_${BACKEND_NAME}, model_non_zero_2d_bool)
+{
+    const auto fn = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/non_zero_2d_bool.prototxt"));
+
+    test_non_zero_constant_folding(fn, {0, 1, 1, 0});
+}

test/type_prop/non_zero.cpp

+//*****************************************************************************
+// Copyright 2017-2020 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#include "gtest/gtest.h"
+#include "ngraph/ngraph.hpp"
+#include "util/type_prop.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+TEST(type_prop, non_zero)
+{
+    auto data = make_shared<op::Parameter>(element::f32, Shape{3, 3, 224, 224});
+    auto non_zero = make_shared<op::NonZero>(data);
+    EXPECT_EQ(non_zero->get_element_type(), element::i64);
+    EXPECT_TRUE(
+        non_zero->get_output_partial_shape(0).same_scheme(PartialShape{4, Dimension::dynamic()}));
+}
+
+TEST(type_prop, non_zero_dynamic)
+{
+    auto data = make_shared<op::Parameter>(element::f32, PartialShape::dynamic());
+    auto non_zero = make_shared<op::NonZero>(data);
+    EXPECT_EQ(non_zero->get_element_type(), element::i64);
+    EXPECT_TRUE(non_zero->get_output_partial_shape(0).same_scheme(
+        PartialShape{Dimension::dynamic(), Dimension::dynamic()}));
+}