return ANEURALNETWORKS_NO_ERROR;
}
+// squeeze(shape) eliminates all the dimensions whose dimensionality is 1
+// For example, squeeze([3, 1, 3]) returns [3, 3]
+static nnfw::util::tensor::Shape squeeze(const nnfw::util::tensor::Shape &shape)
+{
+ nnfw::util::tensor::Shape res(0);
+
+ for (uint32_t axis = 0; axis < shape.rank(); ++axis)
+ {
+ if (shape.dim(axis) != 1)
+ {
+ res.append(shape.dim(axis));
+ }
+ }
+
+ return res;
+}
+
int ANeuralNetworksExecution_setOutput(ANeuralNetworksExecution *execution, int32_t index,
const ANeuralNetworksOperandType *type, void *buffer,
size_t length)
const auto operand_index = execution->plan().model().outputs.at(index);
int32_t output_type = operands.at(operand_index).type();
- if (operands.at(operand_index).shape().rank() == 1)
- {
- const auto len = operands.at(operand_index).shape().dim(0);
+ const auto squeezed_shape = squeeze(operands.at(operand_index).shape());
- asVectorSink(execution, output_type, index, len, buffer, length);
- }
- else if ((operands.at(operand_index).shape().rank() == 2) &&
- (operands.at(operand_index).shape().dim(0) == 1))
+ if (squeezed_shape.rank() == 1)
{
- const auto len = operands.at(operand_index).shape().dim(1);
+ const auto len = squeezed_shape.dim(0);
asVectorSink(execution, output_type, index, len, buffer, length);
}
- else if (operands.at(operand_index).shape().rank() == 4)
+ else if (squeezed_shape.rank() == 3)
{
const auto &operand_shape = operands.at(operand_index).shape().asFeature();