• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2020 Huawei Technologies Co., Ltd
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 #ifndef MINDSPORE_CORE_OPS_SQUEEZE_H_
18 #define MINDSPORE_CORE_OPS_SQUEEZE_H_
19 
20 #include <map>
21 #include <vector>
22 #include <string>
23 #include <memory>
24 #include <algorithm>
25 #include "ops/op_utils.h"
26 #include "ops/primitive_c.h"
27 #include "abstract/primitive_infer_map.h"
28 #include "abstract/abstract_value.h"
29 #include "utils/check_convert_utils.h"
30 
31 namespace mindspore {
32 namespace ops {
33 constexpr auto kNameSqueeze = "Squeeze";
34 /// \brief Returns a tensor with the same data type but dimensions of 1 are removed based on axis.
35 /// Refer to Python API @ref mindspore.ops.Squeeze for more details.
36 class MS_CORE_API Squeeze : public PrimitiveC {
37  public:
38   /// \brief Constructor.
Squeeze()39   Squeeze() : PrimitiveC(kNameSqueeze) { InitIOName({"x"}, {"output"}); }
40   /// \brief Destructor.
41   ~Squeeze() = default;
42   MS_DECLARE_PARENT(Squeeze, PrimitiveC);
43   /// \brief Init. Refer to the parameters of Python API @ref mindspore.ops.Squeeze for the inputs.
44   void Init(const std::vector<int64_t> &axis = {});
45   /// \brief Set axis.
46   void set_axis(const std::vector<int64_t> &axis);
47   /// \brief Get axis.
48   ///
49   /// \return axis.
50   std::vector<int64_t> get_axis() const;
51 };
52 
53 AbstractBasePtr SqueezeInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
54                              const std::vector<AbstractBasePtr> &input_args);
55 using PrimSqueezePtr = std::shared_ptr<Squeeze>;
56 }  // namespace ops
57 }  // namespace mindspore
58 
59 #endif  // MINDSPORE_CORE_OPS_SQUEEZE_H_
60