MultivariateLinearRegressor problem training

Hi everyone, I attempted to use the MultivariateLinearRegressor from the Create ML Components framework to fit some multi-dimensional data linearly (4 dimensions in my example). I aim to obtain multi-dimensional output points (2 points in my example). However, when I fit the model with my training data and test it, it appears that only the first element of my training data is used for training, regardless of whether I use CreateMLComponents.AnnotatedBatch or [CreateMLComponents.AnnotatedFeature, CoreML.MLShapedArray>] as input.

            let sourceMatrix: [[Double]] = [
                [0,0.1,0.2,0.3],
                [0.5,0.2,0.6,0.2]
                
            ]
            
            let referenceMatrix: [[Double]] = [
                [0.2,0.7],
                [0.9,0.1]
            ]

Here is a test code to test the function (ios 18.0 beta, Xcode 16.0 beta) In this example I train the model to learn 2 multidimensional points (4 dimensions) and here are the results of the predictions:

▿ 2 elements
  ▿ 0 : AnnotatedPrediction<MLShapedArray<Double>, MLShapedArray<Double>>
    ▿ prediction : 0.20000000298023224 0.699999988079071 
      ▿ _storage : <StandardStorage<Double>: 0x600002ad8270>
    ▿ annotation : 0.2 0.7 
      ▿ _storage : <StandardStorage<Double>: 0x600002b30600>
  ▿ 1 : AnnotatedPrediction<MLShapedArray<Double>, MLShapedArray<Double>>
    ▿ prediction : 0.23158159852027893 0.9509953260421753 
      ▿ _storage : <StandardStorage<Double>: 0x600002ad8c90>
    ▿ annotation : 0.9 0.1 
      ▿ _storage : <StandardStorage<Double>: 0x600002b55f20>

0.23158159852027893 0.9509953260421753 is totally random and should be far more closer to [0.9,0.1].

Here is the test code : ( i run it on "My mac, Designed for Ipad")

import SwiftUI
import CoreImage
import CoreImage.CIFilterBuiltins
import UIKit
import CoreGraphics
import Accelerate
import Foundation
import CoreML
import CreateML
import CreateMLComponents


func createMLShapedArray(from array: [Double], shape: [Int]) -> MLShapedArray {
    return MLShapedArray(scalars: array, shape: shape)
}


func calculateTransformationMatrixWithNonlinearity(sourceRGB: [[Double]], referenceRGB: [[Double]], degree: Int = 3) async throws -> MultivariateLinearRegressor.Model {

    
    
    
   let annotatedFeatures2 = zip(sourceRGB, referenceRGB).map { (featureArray, targetArray) -> AnnotatedFeature, MLShapedArray> in
        let featureMLShapedArray = createMLShapedArray(from: featureArray, shape: [featureArray.count])
        let targetMLShapedArray = createMLShapedArray(from: targetArray, shape: [targetArray.count])
        return AnnotatedFeature(feature: featureMLShapedArray, annotation: targetMLShapedArray)
    }

    // Flatten the sourceRGBPoly into a single-dimensional array
    var flattenedArray = sourceRGB.flatMap { $0 }
    let featuresMLShapedArray = createMLShapedArray(from: flattenedArray, shape: [2, 4])
    flattenedArray = referenceRGB.flatMap { $0 }
    let targetMLShapedArray = createMLShapedArray(from: flattenedArray, shape: [2, 2])

    // Create AnnotatedFeature instances
   /* let annotatedFeatures2: [AnnotatedFeature, MLShapedArray>] = [
        AnnotatedFeature(feature: featuresMLShapedArray, annotation: targetMLShapedArray)
    ]*/
 
    let annotatedBatch = AnnotatedBatch(features: featuresMLShapedArray, annotations: targetMLShapedArray)

    
    var regressor = MultivariateLinearRegressor()
    regressor.configuration.learningRate = 0.1
    regressor.configuration.maximumIterationCount=5000
    regressor.configuration.batchSize=2
    

    let model = try await regressor.fitted(to: annotatedBatch,validateOn: nil)

    
    //var model = try await regressor.fitted(to: annotatedFeatures2)

    // Proceed to prediction once the model is fitted
    let predictions = try await model.prediction(from: annotatedFeatures2)
    // Process or use the predictions
    print(predictions)
    print("Predictions:", predictions)
    return model
}




struct ContentView: View {
  var body: some View {
    VStack {}
    .onAppear {
      Task {
        do {
            let sourceMatrix: [[Double]] = [
                [0,0.1,0.2,0.3],
                [0.5,0.2,0.6,0.2]
                
            ]
            
            let referenceMatrix: [[Double]] = [
                [0.2,0.7],
                [0.9,0.1]
            ]

          let model = try await calculateTransformationMatrixWithNonlinearity(sourceRGB: sourceMatrix, referenceRGB: referenceMatrix, degree: 2
          )
          print("Model fitted successfully:", model)
        } catch {
          print("Error:", error)
        }
      }
    }
  }
}

import CoreImage
import CoreImage.CIFilterBuiltins
import UIKit
import CoreGraphics
import Accelerate
import Foundation
import CoreML
import CreateML
import CreateMLComponents


func createMLShapedArray(from array: [Double], shape: [Int]) -> MLShapedArray<Double> {
    return MLShapedArray<Double>(scalars: array, shape: shape)
}


func calculateTransformationMatrixWithNonlinearity(sourceRGB: [[Double]], referenceRGB: [[Double]], degree: Int = 3) async throws -> MultivariateLinearRegressor<Double>.Model {

    
    
    
   let annotatedFeatures2 = zip(sourceRGB, referenceRGB).map { (featureArray, targetArray) -> AnnotatedFeature<MLShapedArray<Double>, MLShapedArray<Double>> in
        let featureMLShapedArray = createMLShapedArray(from: featureArray, shape: [featureArray.count])
        let targetMLShapedArray = createMLShapedArray(from: targetArray, shape: [targetArray.count])
        return AnnotatedFeature(feature: featureMLShapedArray, annotation: targetMLShapedArray)
    }

    // Flatten the sourceRGBPoly into a single-dimensional array
    var flattenedArray = sourceRGB.flatMap { $0 }
    let featuresMLShapedArray = createMLShapedArray(from: flattenedArray, shape: [2, 4])
    flattenedArray = referenceRGB.flatMap { $0 }
    let targetMLShapedArray = createMLShapedArray(from: flattenedArray, shape: [2, 2])

    // Create AnnotatedFeature instances
   /* let annotatedFeatures2: [AnnotatedFeature<MLShapedArray<Double>, MLShapedArray<Double>>] = [
        AnnotatedFeature(feature: featuresMLShapedArray, annotation: targetMLShapedArray)
    ]*/
 
    let annotatedBatch = AnnotatedBatch(features: featuresMLShapedArray, annotations: targetMLShapedArray)

    
    var regressor = MultivariateLinearRegressor<Double>()
    regressor.configuration.learningRate = 0.1
    regressor.configuration.maximumIterationCount=5000
    regressor.configuration.batchSize=2
    

    let model = try await regressor.fitted(to: annotatedBatch,validateOn: nil)

    
    //var model = try await regressor.fitted(to: annotatedFeatures2)

    // Proceed to prediction once the model is fitted
    let predictions = try await model.prediction(from: annotatedFeatures2)
    // Process or use the predictions
    print(predictions)
    print("Predictions:", predictions)
    return model
}




struct ContentView: View {
  var body: some View {
    VStack {}
    .onAppear {
      Task {
        do {
            let sourceMatrix: [[Double]] = [
                [0,0.1,0.2,0.3],
                [0.5,0.2,0.6,0.2]
                
            ]
            
            let referenceMatrix: [[Double]] = [
                [0.2,0.7],
                [0.9,0.1]
            ]

          let model = try await calculateTransformationMatrixWithNonlinearity(sourceRGB: sourceMatrix, referenceRGB: referenceMatrix, degree: 2
          )
          print("Model fitted successfully:", model)
        } catch {
          print("Error:", error)
        }
      }
    }
  }
}