Abstract: Certain aspects of the present disclosure provide techniques for partitioning feature maps to improve machine learning model processing. In one aspect, a method, includes partitioning a feature map row-wise into a plurality of feature sub-maps such that: each respective feature sub-map of the plurality of feature sub-maps is defined with respect to a split row determined based on a dense data element count for each row of the feature map; and each feature sub-map of the plurality of feature sub-maps has a same column dimensionality as the feature map; and assigning each of the plurality of feature sub-maps to one of a plurality of tensor compute units and one of a plurality of tensor feature map memory units for processing in parallel.