| 82 | return results; |
| 83 | } |
| 84 | SmallVector<TensorPtr> apply_on_physical_tensor( |
| 85 | const OpDef& def, const SmallVector<TensorPtr>& inputs, |
| 86 | SmallVector<LogicalTensorDesc>& output_descs, const bool& validated) { |
| 87 | auto&& op = def.cast_final_safe<MeshGrid>(); |
| 88 | TensorShape tshp; |
| 89 | TensorShape view_shp; |
| 90 | tshp.ndim = inputs.size(); |
| 91 | view_shp.ndim = inputs.size(); |
| 92 | std::vector<size_t> indexs(inputs.size()); |
| 93 | std::iota(indexs.begin(), indexs.end(), 0); |
| 94 | |
| 95 | if (op.indexing == "xy") { |
| 96 | if (indexs.size() >= 2) { |
| 97 | std::swap(indexs[0], indexs[1]); |
| 98 | } |
| 99 | } else { |
| 100 | mgb_assert(op.indexing == "ij", "meshgrid only support \"ij\" or \"xy\""); |
| 101 | } |
| 102 | for (size_t ind = 0; ind < inputs.size(); ind++) { |
| 103 | auto&& inp = inputs[indexs[ind]]; |
| 104 | mgb_assert(inp->layout().ndim <= 1); |
| 105 | tshp[ind] = inp->layout().total_nr_elems(); |
| 106 | view_shp[ind] = 1; |
| 107 | } |
| 108 | SmallVector<TensorPtr> grids; |
| 109 | for (size_t i = 0; i < inputs.size(); i++) { |
| 110 | auto&& src = inputs[i]; |
| 111 | TensorLayout layout; |
| 112 | view_shp[indexs[i]] = src->layout().total_nr_elems(); |
| 113 | mgb_assert(src->layout().try_reshape(layout, view_shp)); |
| 114 | layout = layout.broadcast(tshp); |
| 115 | view_shp[indexs[i]] = 1; |
| 116 | grids.push_back(Tensor::make(src->blob(), src->offset(), layout)); |
| 117 | } |
| 118 | return grids; |
| 119 | } |
| 120 | OP_TRAIT_REG(MeshGrid, MeshGrid) |
| 121 | .apply_on_var_node(apply_on_var_node) |
| 122 | .infer_output_attrs_fallible(infer_output_attrs_fallible) |
nothing calls this directly
no test coverage detected