ngraph/core/include/ngraph/op/select.hpp

   1 //*****************************************************************************
   2 // Copyright 2017-2020 Intel Corporation
   3 //
   4 // Licensed under the Apache License, Version 2.0 (the "License");
   5 // you may not use this file except in compliance with the License.
   6 // You may obtain a copy of the License at
   7 //
   8 //     http://www.apache.org/licenses/LICENSE-2.0
   9 //
  10 // Unless required by applicable law or agreed to in writing, software
  11 // distributed under the License is distributed on an "AS IS" BASIS,
  12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 // See the License for the specific language governing permissions and
  14 // limitations under the License.
  15 //*****************************************************************************
  16
  17 #pragma once
  18
  19 #include "ngraph/op/op.hpp"
  20
  21 namespace ngraph
  22 {
  23     namespace op
  24     {
  25         namespace v0
  26         {
  27             // clang-format off
  28             /// \brief Elementwise selection operation.
  29             ///
  30             /// ## Inputs
  31             ///
  32             /// |        | Type                                          | Description                                                  |
  33             /// | ------ | --------------------------------------------- | ------------------------------------------------------------ |
  34             /// | `arg0` | \f$\texttt{bool}[d_1,\dots,d_n]~(n \geq 0)\f$ | A tensor of any shape, with element `bool`.                  |
  35             /// | `arg1` | \f$E[d_1,\dots,d_n]~(n \geq 0)\f$             | A tensor of the same shape as `arg0`, with any element type. |
  36             /// | `arg2` | \f$E[d_1,\dots,d_n]~(n \geq 0)\f$             | A tensor of the same shape and element type as `arg1`.       |
  37             ///
  38             /// ## Output
  39             ///
  40             /// | Type                   | Description                                                                                                                                                             |
  41             /// | ---------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
  42             /// | \f$E[d_1,\dots,d_n]\f$ | The tensor \f$T\f$, where \f$T[i_1,\dots,i_n] = \texttt{arg1}[i_1,\dots,i_n]\text{ if }\texttt{arg0}[i_1,\dots,i_n] \neq 0\text{, else }\texttt{arg2}[i_1,\dots,i_n]\f$ |
  43             // clang-format on
  44             class NGRAPH_API Select : public Op
  45             {
  46             public:
  47                 static constexpr NodeTypeInfo type_info{"Select", 0};
  48                 const NodeTypeInfo& get_type_info() const override { return type_info; }
  49                 /// \brief Constructs a selection operation.
  50                 Select() = default;
  51                 /// \brief Constructs a selection operation.
  52                 ///
  53                 /// \param arg0 Node that produces the first input tensor.
  54                 /// \param arg1 Node that produces the second input tensor.
  55                 /// \param arg2 Node that produces the third input tensor.
  56                 Select(const Output<Node>& arg0,
  57                        const Output<Node>& arg1,
  58                        const Output<Node>& arg2);
  59
  60                 virtual std::shared_ptr<Node>
  61                     clone_with_new_inputs(const OutputVector& new_args) const override;
  62                 void validate_and_infer_types() override;
  63             };
  64         }
  65
  66         namespace v1
  67         {
  68             // clang-format off
  69             /// \brief Elementwise selection operation.
  70             ///
  71             /// ## Inputs
  72             ///
  73             /// |        | Type                                          | Description                                                  |
  74             /// | ------ | --------------------------------------------- | ------------------------------------------------------------ |
  75             /// | `arg0` | \f$\texttt{bool}[d_1,\dots,d_n]~(n \geq 0)\f$ | A tensor of any shape, with element `bool`.                  |
  76             /// | `arg1` | \f$E[d_1,\dots,d_n]~(n \geq 0)\f$             | A tensor of a shape that is broadcast-compatible with `arg0`, with any element type. |
  77             /// | `arg2` | \f$E[d_1,\dots,d_n]~(n \geq 0)\f$             | A tensor of a shape that is broadcast-compatible with `arg0`, and same element type as `arg1`. |
  78             /// | `auto_broadcast`| AutoBroadcastSpec                             | Auto broadcast specification.                                |
  79             ///
  80             /// ## Output
  81             ///
  82             /// | Type                   | Description                                                                                                                                                             |
  83             /// | ---------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
  84             /// | \f$E[d_1,\dots,d_n]\f$ | The tensor \f$T\f$, where \f$T[i_1,\dots,i_n] = \texttt{arg1}[i_1,\dots,i_n]\text{ if }\texttt{arg0}[i_1,\dots,i_n] \neq 0\text{, else }\texttt{arg2}[i_1,\dots,i_n]\f$ |
  85             // clang-format on
  86             class NGRAPH_API Select : public Op
  87             {
  88             public:
  89                 NGRAPH_RTTI_DECLARATION;
  90                 /// \brief Constructs a selection operation.
  91                 Select()
  92                     : m_auto_broadcast(AutoBroadcastSpec(AutoBroadcastType::NUMPY))
  93                 {
  94                 }
  95
  96                 /// \brief Constructs a selection operation.
  97                 ///
  98                 /// \param arg0 Node that produces the first input tensor.
  99                 /// \param arg1 Node that produces the second input tensor.
 100                 /// \param arg2 Node that produces the third input tensor.
 101                 /// \param auto_broadcast Auto broadcast specification. Default is Numpy-style
 102                 ///                       implicit broadcasting.
 103                 Select(const Output<Node>& arg0,
 104                        const Output<Node>& arg1,
 105                        const Output<Node>& arg2,
 106                        const AutoBroadcastSpec& auto_broadcast =
 107                            AutoBroadcastSpec(AutoBroadcastType::NUMPY));
 108
 109                 virtual std::shared_ptr<Node>
 110                     clone_with_new_inputs(const OutputVector& new_args) const override;
 111                 bool visit_attributes(AttributeVisitor& visitor) override;
 112                 void validate_and_infer_types() override;
 113
 114                 const AutoBroadcastSpec& get_auto_broadcast() const { return m_auto_broadcast; }
 115                 void set_auto_broadcast(const AutoBroadcastSpec& auto_broadcast)
 116                 {
 117                     m_auto_broadcast = auto_broadcast;
 118                 }
 119                 // TODO: Move all uses of get_autob to get_auto_broadcast() and remove this.
 120                 const AutoBroadcastSpec& get_autob() const override { return m_auto_broadcast; }
 121             private:
 122                 AutoBroadcastSpec m_auto_broadcast;
 123             };
 124         }
 125         using v0::Select;
 126     }
 127 }