How to reshape MLMultipArray in Swift

I'm trying to colorize image using coreML, and the model always returns (1, 3, 256, 256) as result. I'm reshaping to convert it to (3, 256, 256) but not sure if it's correct way. Is there no way to pick only first element in array?

let input = try colorizer_coreInput(image_smallWith: resized)
let predicted = try colorizer_core(configuration: MLModelConfiguration()).prediction(input: input)
let array = predicted.colorized_small_image


var colori: UIImage?
try withExtendedLifetime(array) {
   if let reshapedArray = try? array.reshaped(to: [3, 256, 256]) {
       colori = createUIImage(fromFloatArray: reshapedArray, min: 0, max: 255)
   }
}

Solution

I think the following extension will provide what you need. It's not a general reshape, but for your specific case of dropping a first dimension of 1, it will do the job. It could be generalized to drop any of the dimensions so long as the dropped dimension size is 1, because that doesn't change the underlying memory buffer. The moment you drop a dimension whose size is larger than 1, things are more complicated.

Note there are two versions of the method. The first one is fast, but unsafe if the source MLMultiArray goes out of scope, but since in your example it would be used in the scope of a withExtendedLifetime it's the one I'd use.

The second version is safe, but somewhat slower because it has to copy data into a newly allocated buffer.

extension MLMultiArray
{
    /**
     This version simply re-uses the existing `MLMultiArray`'s buffer, as a
     result  this version is very fast; however, if the original `MLMultiArray`
     is deinited this the new `MLMultiArray`'s buffer will be invalid.
     Therefore, ONLY use this when you can keep the original `MLMultiArray`
     alive for as long the one returned by this method is needed.
     */
    func dropFirstDimensionNoCopy() -> MLMultiArray?
    {
        guard shape[0] == 1 else
        {
            assertionFailure("Dropping non-1 dimension not supported")
            return nil
        }
        
        return try? MLMultiArray(
            dataPointer: self.dataPointer,
            shape: [NSNumber](self.shape.dropFirst()),
            dataType: self.dataType,
            strides: [NSNumber](self.strides.dropFirst()),
            deallocator: { _ in }
        )
    }
    
    func dropFirstDimension() -> MLMultiArray?
    {
        guard shape[0] == 1 else
        {
            assertionFailure("Dropping non-1 dimension not supported")
            return nil
        }
        
        var elementSize: Int
        switch dataType
        {
            case .double  : elementSize = 8
                
            case .float,
                 .float32,
                 .int32   : elementSize = 4
                
            case .float16 : elementSize = 2
                
            @unknown default:
                fatalError("Unknown MLMultiArrayDataType, \(dataType)")
        }
        
        let bufferSize = strides.first!.intValue * elementSize
        
        guard let newArray = try? MLMultiArray(
            shape: .init(shape.dropFirst()),
            dataType: dataType)
        else { fatalError("Unable to create MLMultiArray.") }
        
        memcpy(newArray.dataPointer, dataPointer, bufferSize)
        return newArray
    }
}