tensor (sequence) mask (#735)

* tensor (sequence) mask

* convert sequence to match input type

* update API to take batch axis

src/ngraph/CMakeLists.txt

@@ -19,6 +19,7 @@ set (SRC
     builder/autobroadcast.cpp
     builder/numpy_transpose.cpp
     builder/reduce_ops.cpp
+    builder/tensor_mask.cpp
     coordinate_transform.cpp
     descriptor/input.cpp
     descriptor/layout/dense_tensor_view_layout.cpp

src/ngraph/builder/tensor_mask.cpp

+/*******************************************************************************
+* Copyright 2017-2018 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/builder/tensor_mask.hpp"
+
+#include "ngraph/op/broadcast.hpp"
+#include "ngraph/op/constant.hpp"
+#include "ngraph/op/convert.hpp"
+#include "ngraph/op/less.hpp"
+#include "ngraph/op/reshape.hpp"
+
+using namespace ngraph;
+
+std::shared_ptr<Node> ngraph::builder::tensor_mask(const std::shared_ptr<Node>& sequence_lengths,
+                                                   size_t sequence_axis,
+                                                   size_t batch_axis,
+                                                   Shape mask_shape)
+{
+    if (sequence_axis >= mask_shape.size())
+    {
+        throw ngraph_error("Sequence axis must be in range 0..mask_shape rank");
+    }
+
+    if (batch_axis >= mask_shape.size())
+    {
+        throw ngraph_error("Sequence axis must be in range 0..mask_shape rank");
+    }
+
+    // all axes except the sequence axis
+    AxisSet non_sequence_axes;
+    // all axes except the batch axis
+    AxisSet non_batch_axes;
+
+    for (auto axis = 0; axis < mask_shape.size(); ++axis)
+    {
+        if (axis != sequence_axis)
+        {
+            non_sequence_axes.insert(axis);
+        }
+        if (axis != batch_axis)
+        {
+            non_batch_axes.insert(axis);
+        }
+    }
+
+    // broadcast sequence lengths to mask shape along all non-batch axes
+    auto broadcast_sequence_lengths =
+        std::make_shared<op::Broadcast>(sequence_lengths, mask_shape, non_batch_axes);
+
+    // create sequence data [0, ..., max_sequence_length]
+    auto max_sequence_length = mask_shape[sequence_axis];
+    std::vector<uint32_t> sequence_data(max_sequence_length);
+    std::iota(sequence_data.begin(), sequence_data.end(), 0);
+
+    // create sequence constant
+    auto sequence =
+        std::make_shared<op::Constant>(element::u32, Shape{max_sequence_length}, sequence_data);
+
+    // convert sequence to input type
+    auto convert_sequence =
+        std::make_shared<op::Convert>(sequence, sequence_lengths->get_element_type());
+
+    // broadcast sequence to mask shape along all non-sequence axes
+    auto broadcast_sequence =
+        std::make_shared<op::Broadcast>(convert_sequence, mask_shape, non_sequence_axes);
+
+    // mask = sequence_length < sequence
+    return std::make_shared<op::Less>(broadcast_sequence, broadcast_sequence_lengths);
+}

src/ngraph/builder/tensor_mask.hpp

+/*******************************************************************************
+* Copyright 2017-2018 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/axis_set.hpp"
+#include "ngraph/node.hpp"
+#include "ngraph/shape.hpp"
+
+namespace ngraph
+{
+    namespace builder
+    {
+        std::shared_ptr<Node> tensor_mask(const std::shared_ptr<Node>& sequence_lengths,
+                                          size_t sequence_axis,
+                                          size_t batch_axis,
+                                          Shape mask_shape);
+    }
+}
\ No newline at end of file

src/ngraph/ngraph.hpp

@@ -46,6 +46,7 @@
 #include "ngraph/builder/autobroadcast.hpp"
 #include "ngraph/builder/numpy_transpose.hpp"
 #include "ngraph/builder/reduce_ops.hpp"
+#include "ngraph/builder/tensor_mask.hpp"
 #include "ngraph/coordinate_transform.hpp"
 #include "ngraph/descriptor/buffer.hpp"
 #include "ngraph/descriptor/input.hpp"

test/builder.cpp

@@ -140,3 +140,29 @@ TEST(builder, numpy_transpose)
     EXPECT_ANY_THROW(
         dynamic_pointer_cast<op::Reshape>(builder::numpy_transpose(param, AxisVector{2, 2, 1})));
 }
+
+TEST(builder, tensor_mask)
+{
+    Shape max_sequence_length{3};
+    auto sequence_lengths = make_shared<op::Parameter>(element::u32, max_sequence_length);
+
+    Shape mask_shape{3, 5};
+    auto f =
+        make_shared<Function>(builder::tensor_mask(sequence_lengths, 1, 0, mask_shape),
+                              op::ParameterVector{sequence_lengths});
+
+    auto manager = runtime::Manager::get("INTERPRETER");
+    auto external = manager->compile(f);
+    auto backend = manager->allocate_backend();
+    auto cf = backend->make_call_frame(external);
+
+    auto sequence_lengths_data =
+        backend->make_primary_tensor_view(element::u32, max_sequence_length);
+    copy_data(sequence_lengths_data, vector<uint32_t>{1, 3, 2});
+    auto result = backend->make_primary_tensor_view(element::boolean, mask_shape);
+
+    cf->call({result}, {sequence_lengths_data});
+    vector<char> expected{1, 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0};
+
+    EXPECT_EQ(expected, read_vector<char>(result));
+}