Optimized MLMultiArray Access

I’m looking for feedback on a more efficient way to access/filter an MLMultiArray. I’ve looked around and seen sites imply that I can use map or compactmap but I have been able to make those work.

My objective: Check the 5th element on the second dimension (my confidence score) and if it is greater than my threshold get the values from that dimension and store them as a valid detection.

Right now it feels like I’m brute forcing it. Any suggestions are appreciated:

        var filtered_results = [[Float]]()
        let confidence_threshold:Float = 0.5
        var temp = [Float]()
        key = [0,4,0]
        for j in 0...((prediction_multi?.shape[2].intValue)!-1){
            key = [0,4,j] as [NSNumber]
            temp = []
            if (prediction_multi?[key].floatValue)! > confidence_threshold {
                for i in 0...((prediction_multi?.shape[1].intValue)!-1){
                    key = [0,i,j] as [NSNumber]
                    temp.append((prediction_multi?[key].floatValue)!)
                }
                filtered_results.append(temp)
            }
        }

It works and does exactly what I want, but it just feels* wrong and complex. In Python I just do:

if predictions[0,4,i]> confidence_threshold:
final_results.append(predictions[0,:,i])

replying to my own thread…Found this solution.

        extension Collection where Self.Iterator.Element: RandomAccessCollection {
            // PRECONDITION: `self` must be rectangular, i.e. every row has equal size.
            func transposed() -> [[Self.Iterator.Element.Iterator.Element]] {
                guard let firstRow = self.first else { return [] }
                return firstRow.indices.map { index in
                    self.map{ $0[index] }
                }
            }
        }

        var test = [[1,2,3],[1,2,3],[1,2,3],[1,2,3]]
        var results = test.transposed()
        print(results) // [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]]