Home
last modified time | relevance | path

Searched refs:axis_tensor (Results 1 – 12 of 12) sorted by relevance

/third_party/mindspore/mindspore/core/ops/
Dreduce_sum.cc107 auto axis_tensor = input_args[1]->cast<abstract::AbstractTensorPtr>(); in InferShape() local
108 auto axis_shape = axis_tensor->shape()->shape(); in InferShape()
144 auto axis_tensor = axis_ptr->cast<tensor::TensorPtr>(); in InferShape() local
145 MS_EXCEPTION_IF_NULL(axis_tensor); in InferShape()
146 size_t data_size = LongToSize(axis_tensor->DataSize()); in InferShape()
149 auto shape_data = reinterpret_cast<int *>(axis_tensor->data_c()); in InferShape()
156 auto shape_data2 = reinterpret_cast<int64_t *>(axis_tensor->data_c()); in InferShape()
Dgather.cc56 auto axis_tensor = axis_value_ptr->cast<tensor::TensorPtr>(); in GatherInfer() local
57 MS_EXCEPTION_IF_NULL(axis_tensor); in GatherInfer()
58 axis_val = *static_cast<int64_t *>(axis_tensor->data_c()); in GatherInfer()
/third_party/mindspore/mindspore/lite/src/ops/compat/v0/
Dgather_compat_v0.cc37 auto axis_tensor = AttrToTensor(&axis_attr, 1, false, kNumberTypeInt32, tensor_bufs); in TransferGatherAttr() local
38 if (axis_tensor == nullptr) { in TransferGatherAttr()
42 dst_tensors->push_back(axis_tensor); in TransferGatherAttr()
/third_party/mindspore/mindspore/lite/src/delegate/tensorrt/op/
Dreduce_tensorrt.cc131 mindspore::MSTensor axis_tensor = this->in_tensors_[1]; in GetAxis() local
132 if (axis_tensor.Data() == nullptr) { in GetAxis()
136 if (axis_tensor.DataType() != DataType::kNumberTypeInt32) { in GetAxis()
139 int *axis_data = reinterpret_cast<int *>(axis_tensor.MutableData()); in GetAxis()
140 for (int i = 0; i < axis_tensor.ElementNum(); i++) { in GetAxis()
/third_party/mindspore/mindspore/lite/tools/converter/parser/tflite/
Dtflite_argmin_parser.cc38 const auto &axis_tensor = tflite_subgraph->tensors.at(tflite_op->inputs[1]); in Parse() local
39 MS_CHECK_TRUE_MSG(axis_tensor != nullptr, nullptr, "axis_tensor is nullptr"); in Parse()
40 const auto &buf_data = tflite_model->buffers.at(axis_tensor->buffer); in Parse()
Dtflite_argmax_parser.cc38 const auto &axis_tensor = tflite_subgraph->tensors.at(tflite_op->inputs[1]); in Parse() local
39 MS_CHECK_TRUE_MSG(axis_tensor != nullptr, nullptr, "axis_tensor is nullptr"); in Parse()
40 const auto &buf_data = tflite_model->buffers.at(axis_tensor->buffer); in Parse()
Dtflite_split_v_parser.cc56 const auto &axis_tensor = tflite_subgraph->tensors.at(tflite_op->inputs.at(2)); in Parse() local
57 if (axis_tensor == nullptr) { in Parse()
61 auto &axis_buf_data = tflite_model->buffers.at(axis_tensor->buffer); in Parse()
Dtflite_split_parser.cc48 const auto &axis_tensor = tflite_subgraph->tensors.at(tflite_op->inputs.at(0)); in Parse() local
49 if (axis_tensor == nullptr) { in Parse()
53 auto &axis_buf_data = tflite_model->buffers.at(axis_tensor->buffer); in Parse()
/third_party/mindspore/mindspore/lite/tools/converter/parser/tf/
Dtf_argmin_parser.cc38 auto &axis_tensor = attr_value.tensor(); in Parse() local
39 prim->set_axis(axis_tensor.int_val(0)); in Parse()
Dtf_argmax_parser.cc43 auto &axis_tensor = attr_value.tensor(); in Parse() local
44 prim->set_axis(axis_tensor.int_val(0)); in Parse()
/third_party/mindspore/mindspore/lite/src/runtime/kernel/arm/fp32/
Dcumsum_fp32.cc63 auto axis_tensor = in_tensors_.at(1); in ReSize() local
64 CHECK_NULL_RETURN(axis_tensor); in ReSize()
65 int *axis_data = reinterpret_cast<int *>(axis_tensor->data()); in ReSize()
/third_party/mindspore/mindspore/core/abstract/
Dprim_arrays.cc596 auto axis_tensor = axis_value_ptr->cast<tensor::TensorPtr>(); in InferImplGatherV2() local
597 MS_EXCEPTION_IF_NULL(axis_tensor); in InferImplGatherV2()
598 axis_val = *static_cast<int64_t *>(axis_tensor->data_c()); in InferImplGatherV2()