[Colab] The number of scalars does not match the shape MNIST example

SwiftAI
1

I’m trying to load the MNIST dataset as in the tensorflow repo https://github.com/tensorflow/swift-models/blob/master/MNIST/main.swift

Here is the definition of the helper functions used later

func download(_ url: String, _ PATH: Path = Path.cwd) -> String {
    let fname = url.split(separator:"/").last!
    let path = (PATH/fname).string
    let r = Just.get(url)
    if r.ok {
        let body = r.content!
        do {
            try body.write(to: URL(fileURLWithPath: path))
        } catch {
            print("Failed to write \(url) due to error \(error)")
        }
    }
    return path
}
/// Reads a file into an array of bytes.
func readFile(_ path: String) -> [UInt8] {
    let data = try! Data(contentsOf: URL(fileURLWithPath: path), options: [])
    return [UInt8](data)
}

I download the two files like this

> let base_url = "https://github.com/tensorflow/swift-models/raw/master/MNIST"
> let train_images_idx3_ubyte = download("\(base_url)/train-images-idx3-ubyte", PATH)
> let train_labels_idx3_ubyte = download("\(base_url)/train-labels-idx1-ubyte", PATH)

I read the downloaded files

> let images = readFile(train_images_idx3_ubyte).dropFirst(16).map(Float.init)
> let labels = readFile(train_labels_idx3_ubyte).dropFirst(8).map(Int32.init)

> let rowCount = labels.count
> let imageHeight = 28, imageWidth = 28

Now when I try to load the raw data into tensors I got the error

> Tensor(shape: [rowCount, 1, imageHeight, imageWidth], scalars: images)
Precondition failed: The number of scalars does not match the shape.: file /swift-base/swift/stdlib/public/TensorFlow/Tensor.swift, line 101
Current stack trace:
0    libswiftCore.so                    0x00007f1994e864a0 _swift_stdlib_reportFatalErrorInFile + 115
1    libswiftCore.so                    0x00007f1994dce30c <unavailable> + 3035916
2    libswiftCore.so                    0x00007f1994dce3fe <unavailable> + 3036158
3    libswiftCore.so                    0x00007f1994c156c2 <unavailable> + 1230530
4    libswiftCore.so                    0x00007f1994d9b292 <unavailable> + 2826898
5    libswiftCore.so                    0x00007f1994c14ba9 <unavailable> + 1227689
6    libswiftTensorFlow.so              0x00007f19913e74d0 __tf_tensor_from_scalars + 496
Current stack trace:
	frame #4: 0x00007f19c2494985 $__lldb_expr144`main at <Cell 23>:1
	frame #15: 0x00007f19c23f5129 $__lldb_expr122`main at <Cell 20>:1

Any idea why I’m the shape does not match the expected rowCount*28*28= 47165440?