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 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)
        }
      }
    }
  }
}