| 36 | std::vector<op_desc> operators() const { return {{"ExpandDims"}}; } |
| 37 | |
| 38 | instruction_ref parse(const op_desc& /*opd*/, |
| 39 | const tf_parser& /*parser*/, |
| 40 | const tf_parser::node_info& info, |
| 41 | std::vector<instruction_ref> args) const |
| 42 | { |
| 43 | std::vector<size_t> input_dims = args[0]->get_shape().lens(); |
| 44 | std::vector<int64_t> new_dims(input_dims.begin(), input_dims.end()); |
| 45 | size_t num_dims = input_dims.size(); |
| 46 | int32_t dim = args[1]->eval().at<int32_t>(); |
| 47 | |
| 48 | if(dim < 0) |
| 49 | { |
| 50 | new_dims.insert(new_dims.begin() + (num_dims + dim + 1), 1); |
| 51 | } |
| 52 | else |
| 53 | { |
| 54 | new_dims.insert(new_dims.begin() + dim, 1); |
| 55 | } |
| 56 | return info.add_instruction(make_op("reshape", {{"dims", new_dims}}), args[0]); |
| 57 | } |
| 58 | }; |
| 59 | |
| 60 | } // namespace tf |