Function ComputeSplitAxis

src/backends/backendsCommon/WorkloadUtils.cpp:377–400

}

std::set<unsigned int> ComputeSplitAxis(const armnn::SplitterDescriptor& desc, const TensorShape& input)
{
    unsigned int numSplit = desc.GetNumViews();
    unsigned int numDimensions = desc.GetNumDimensions();
    std::set<unsigned int> splitAxis;
    if (desc.HasAxis())
    {
        splitAxis.insert(armnnUtils::GetUnsignedAxis(desc.GetNumDimensions(), desc.GetAxis()));
    }
    else
    {
        for (unsigned int i = 0; i < numSplit; ++i)
        {
            for (unsigned int dimIdx = 0; dimIdx < numDimensions; ++dimIdx)
            {
                if (desc.GetViewSizes(i)[dimIdx] != input[dimIdx])
                {
                    splitAxis.insert(dimIdx);
                }
            }
        }
    }
    return splitAxis;
}

} // namespace armnn