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_DEPRECATED(
  45                 "This operation is deprecated and will be removed soon. "
  46                 "Use v1::Select instead of it.") NGRAPH_API Select : public Op
  47             {
  48                 NGRAPH_SUPPRESS_DEPRECATED_START
  49             public:
  50                 static constexpr NodeTypeInfo type_info{"Select", 0};
  51                 const NodeTypeInfo& get_type_info() const override { return type_info; }
  52                 /// \brief Constructs a selection operation.
  53                 Select() = default;
  54                 /// \brief Constructs a selection operation.
  55                 ///
  56                 /// \param arg0 Node that produces the first input tensor.
  57                 /// \param arg1 Node that produces the second input tensor.
  58                 /// \param arg2 Node that produces the third input tensor.
  59                 Select(const Output<Node>& arg0,
  60                        const Output<Node>& arg1,
  61                        const Output<Node>& arg2);
  62
  63                 virtual std::shared_ptr<Node>
  64                     clone_with_new_inputs(const OutputVector& new_args) const override;
  65                 void validate_and_infer_types() override;
  66                 NGRAPH_SUPPRESS_DEPRECATED_END
  67             };
  68         }
  69
  70         namespace v1
  71         {
  72             // clang-format off
  73             /// \brief Elementwise selection operation.
  74             ///
  75             /// ## Inputs
  76             ///
  77             /// |        | Type                                          | Description                                                  |
  78             /// | ------ | --------------------------------------------- | ------------------------------------------------------------ |
  79             /// | `arg0` | \f$\texttt{bool}[d_1,\dots,d_n]~(n \geq 0)\f$ | A tensor of any shape, with element `bool`.                  |
  80             /// | `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. |
  81             /// | `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`. |
  82             /// | `auto_broadcast`| AutoBroadcastSpec                             | Auto broadcast specification.                                |
  83             ///
  84             /// ## Output
  85             ///
  86             /// | Type                   | Description                                                                                                                                                             |
  87             /// | ---------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
  88             /// | \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$ |
  89             // clang-format on
  90             class NGRAPH_API Select : public Op
  91             {
  92             public:
  93                 NGRAPH_RTTI_DECLARATION;
  94                 /// \brief Constructs a selection operation.
  95                 Select()
  96                     : m_auto_broadcast(AutoBroadcastSpec(AutoBroadcastType::NUMPY))
  97                 {
  98                 }
  99
 100                 /// \brief Constructs a selection operation.
 101                 ///
 102                 /// \param arg0 Node that produces the first input tensor.
 103                 /// \param arg1 Node that produces the second input tensor.
 104                 /// \param arg2 Node that produces the third input tensor.
 105                 /// \param auto_broadcast Auto broadcast specification. Default is Numpy-style
 106                 ///                       implicit broadcasting.
 107                 Select(const Output<Node>& arg0,
 108                        const Output<Node>& arg1,
 109                        const Output<Node>& arg2,
 110                        const AutoBroadcastSpec& auto_broadcast =
 111                            AutoBroadcastSpec(AutoBroadcastType::NUMPY));
 112
 113                 virtual std::shared_ptr<Node>
 114                     clone_with_new_inputs(const OutputVector& new_args) const override;
 115                 bool visit_attributes(AttributeVisitor& visitor) override;
 116                 void validate_and_infer_types() override;
 117
 118                 const AutoBroadcastSpec& get_auto_broadcast() const { return m_auto_broadcast; }
 119                 void set_auto_broadcast(const AutoBroadcastSpec& auto_broadcast)
 120                 {
 121                     m_auto_broadcast = auto_broadcast;
 122                 }
 123                 // TODO: Move all uses of get_autob to get_auto_broadcast() and remove this.
 124                 const AutoBroadcastSpec& get_autob() const override { return m_auto_broadcast; }
 125             private:
 126                 AutoBroadcastSpec m_auto_broadcast;
 127             };
 128         }
 129         NGRAPH_SUPPRESS_DEPRECATED_START
 130         using v0::Select;
 131         NGRAPH_SUPPRESS_DEPRECATED_END
 132     }
 133 }