Function SizeOutAxis

dnn/src/cambricon/softmax/opr_impl.cpp:24–30

}

size_t SizeOutAxis(const int axis, const size_t* dims, const int ndim) {
    size_t size = 1;
    for (int i = axis + 1; i < ndim; i++) {
        size *= dims[i];
    }
    return size;
}

SmallVector<size_t> init_shape(_megdnn_tensor_in src, const int param_axis) {
    auto dims = src.layout.shape;