diff --git a/Datasets/CIFAR10/CIFAR10.swift b/Datasets/CIFAR10/CIFAR10.swift index 7b74380adbd..3ba5650d0dd 100644 --- a/Datasets/CIFAR10/CIFAR10.swift +++ b/Datasets/CIFAR10/CIFAR10.swift @@ -113,7 +113,7 @@ func loadCIFARFile(named name: String, in directory: String = ".") -> CIFARExamp let images = Tensor(shape: [imageCount, 3, 32, 32], scalars: bytes) // Transpose from the CIFAR-provided N(CHW) to TF's default NHWC. - let imageTensor = Tensor(images.transposed(withPermutations: [0, 2, 3, 1])) + let imageTensor = Tensor(images.transposed(permutation: [0, 2, 3, 1])) let mean = Tensor([0.485, 0.456, 0.406]) let std = Tensor([0.229, 0.224, 0.225]) @@ -125,8 +125,8 @@ func loadCIFARFile(named name: String, in directory: String = ".") -> CIFARExamp func loadCIFARTrainingFiles() -> CIFARExample { let data = (1..<6).map { loadCIFARFile(named: "data_batch_\($0).bin") } return CIFARExample( - label: Raw.concat(concatDim: Tensor(0), data.map { $0.label }), - data: Raw.concat(concatDim: Tensor(0), data.map { $0.data }) + label: _Raw.concat(concatDim: Tensor(0), data.map { $0.label }), + data: _Raw.concat(concatDim: Tensor(0), data.map { $0.data }) ) } diff --git a/Datasets/MNIST/MNIST.swift b/Datasets/MNIST/MNIST.swift index 9e984799dd8..6a79b4b5d65 100644 --- a/Datasets/MNIST/MNIST.swift +++ b/Datasets/MNIST/MNIST.swift @@ -104,7 +104,7 @@ fileprivate func fetchDataset( return ( images: Tensor(shape: [rowCount, 1, imageHeight, imageWidth], scalars: images) - .transposed(withPermutations: [0, 2, 3, 1]) / 255, // NHWC + .transposed(permutation: [0, 2, 3, 1]) / 255, // NHWC labels: Tensor(labels) ) } diff --git a/GAN/main.swift b/GAN/main.swift index f57a550274f..8e09985ec34 100644 --- a/GAN/main.swift +++ b/GAN/main.swift @@ -126,7 +126,7 @@ func saveImageGrid(_ testImage: Tensor, name: String) throws { // Add padding. gridImage = gridImage.padded(forSizes: [(0, 0), (0, 0), (1, 1), (1, 1)], with: 1) // Transpose to create single image. - gridImage = gridImage.transposed(withPermutations: [0, 2, 1, 3]) + gridImage = gridImage.transposed(permutation: [0, 2, 1, 3]) gridImage = gridImage.reshaped( to: [ (imageHeight + 2) * testImageGridSize, diff --git a/MiniGo/Models/PythonCheckpointReader.swift b/MiniGo/Models/PythonCheckpointReader.swift index 2fe1a950032..58c1746a7e8 100644 --- a/MiniGo/Models/PythonCheckpointReader.swift +++ b/MiniGo/Models/PythonCheckpointReader.swift @@ -28,7 +28,7 @@ public class PythonCheckpointReader { let countSuffix = layerCounts[layerName] == nil ? "" : "_\(layerCounts[layerName]!)" let tensorName = layerName + countSuffix + "/" + weightName // TODO(jekbradbury): support variadic dtype attrs in RawOpsGenerated - return Raw.restoreV2(prefix: StringTensor(path), + return _Raw.restoreV2(prefix: StringTensor(path), tensorNames: StringTensor([tensorName]), shapeAndSlices: StringTensor([""])) } diff --git a/Support/Image.swift b/Support/Image.swift index 8cef015a6c0..2df5d537ea4 100644 --- a/Support/Image.swift +++ b/Support/Image.swift @@ -37,11 +37,11 @@ public struct Image { } public init(jpeg url: URL, byteOrdering: ByteOrdering = .rgb) { - let loadedFile = Raw.readFile(filename: StringTensor(url.absoluteString)) - let loadedJpeg = Raw.decodeJpeg(contents: loadedFile, channels: 3, dctMethod: "") + let loadedFile = _Raw.readFile(filename: StringTensor(url.absoluteString)) + let loadedJpeg = _Raw.decodeJpeg(contents: loadedFile, channels: 3, dctMethod: "") if byteOrdering == .bgr { self.imageData = .uint8( - data: Raw.reverse(loadedJpeg, dims: Tensor([false, false, false, true]))) + data: _Raw.reverse(loadedJpeg, dims: Tensor([false, false, false, true]))) } else { self.imageData = .uint8(data: loadedJpeg) } @@ -59,21 +59,21 @@ public struct Image { outputImageData = Tensor(adjustedData) } - let encodedJpeg = Raw.encodeJpeg( + let encodedJpeg = _Raw.encodeJpeg( image: outputImageData, format: .grayscale, quality: quality, xmpMetadata: "") - Raw.writeFile(filename: StringTensor(url.absoluteString), contents: encodedJpeg) + _Raw.writeFile(filename: StringTensor(url.absoluteString), contents: encodedJpeg) } public func resized(to size: (Int, Int)) -> Image { switch self.imageData { case let .uint8(data): return Image( - tensor: Raw.resizeBilinear( + tensor: _Raw.resizeBilinear( images: Tensor([data]), size: Tensor([Int32(size.0), Int32(size.1)]))) case let .float(data): return Image( - tensor: Raw.resizeBilinear( + tensor: _Raw.resizeBilinear( images: Tensor([data]), size: Tensor([Int32(size.0), Int32(size.1)]))) } diff --git a/Transformer/Model.swift b/Transformer/Model.swift index 20018da23a0..cc9181863f8 100644 --- a/Transformer/Model.swift +++ b/Transformer/Model.swift @@ -92,7 +92,7 @@ func causallyMasked(_ dotProducts: Tensor, enable: Bool = false) -> Tenso } let (queryTimeSteps, keyTimeSteps) = (dotProducts.shape[1], dotProducts.shape[2]) let ones = Tensor(ones: [1, queryTimeSteps, keyTimeSteps]) - let mask = Raw.matrixBandPart( + let mask = _Raw.matrixBandPart( ones, numLower: Tensor(Int32(-1)), numUpper: Tensor(Int32(queryTimeSteps - keyTimeSteps))) @@ -138,7 +138,7 @@ func splitHeads(_ input: Tensor, headCount: Int) -> Tensor { let (batchSize, timeSteps, features) = (input.shape[0], input.shape[1], input.shape[2]) let featuresPerHead = features / headCount let splitLastDim = input.reshaped(to: [batchSize, timeSteps, headCount, featuresPerHead]) - let movedToFront = splitLastDim.transposed(withPermutations: 0, 2, 1, 3) + let movedToFront = splitLastDim.transposed(permutation: 0, 2, 1, 3) return movedToFront.reshaped(to: [batchSize * headCount, timeSteps, featuresPerHead]) } @@ -149,7 +149,7 @@ func joinHeads(_ input: Tensor, headCount: Int) -> Tensor { let batchSize = generalizedBatch / headCount let features = featuresPerHead * headCount let splitFirstDim = input.reshaped(to: [batchSize, headCount, timeSteps, featuresPerHead]) - let movedToBack = splitFirstDim.transposed(withPermutations: 0, 2, 1, 3) + let movedToBack = splitFirstDim.transposed(permutation: 0, 2, 1, 3) return movedToBack.reshaped(to: [batchSize, timeSteps, features]) } @@ -173,7 +173,7 @@ func _vjpSplitQKV(_ input: Tensor) -> (AttentionInput, (AttentionInput.TangentVector) -> Tensor) { let value = splitQKV(input) return (value, { seed in - return Raw.concatV2([seed.query, seed.key, seed.value], axis: Tensor(2)) + return _Raw.concatV2([seed.query, seed.key, seed.value], axis: Tensor(2)) }) } diff --git a/Transformer/Operators.swift b/Transformer/Operators.swift index 20944d45bd9..60a5e82b7f5 100644 --- a/Transformer/Operators.swift +++ b/Transformer/Operators.swift @@ -37,7 +37,7 @@ func batchedMatmul( adjointLeft: Bool = false, adjointRight: Bool = false ) -> Tensor { - return Raw.batchMatMul(left, right, adjX: adjointLeft, adjY: adjointRight) + return _Raw.batchMatMul(left, right, adjX: adjointLeft, adjY: adjointRight) } @usableFromInline diff --git a/Transformer/PythonCheckpointReader.swift b/Transformer/PythonCheckpointReader.swift index 3b84c38e1a6..9597cd28e25 100644 --- a/Transformer/PythonCheckpointReader.swift +++ b/Transformer/PythonCheckpointReader.swift @@ -36,7 +36,7 @@ func readTensor( scalarType: Scalar.Type ) -> Tensor { // TODO(jekbradbury): support variadic dtype attrs in RawOpsGenerated - return Raw.restoreV2(prefix: StringTensor(path), + return _Raw.restoreV2(prefix: StringTensor(path), tensorNames: StringTensor([name]), shapeAndSlices: StringTensor([""])) } diff --git a/Transformer/main.swift b/Transformer/main.swift index 0cc91eb2171..a3d4a4fa5f7 100644 --- a/Transformer/main.swift +++ b/Transformer/main.swift @@ -53,7 +53,7 @@ for _ in 0..<100 { let lastLogit = logits.slice( lowerBounds: [0, timeSteps - 1, 0], upperBounds: [batchSize, timeSteps, vocabSize]) / temperature - tokens = Raw.multinomial(logits: lastLogit.squeezingShape(at: 1), numSamples: Tensor(1)) + tokens = _Raw.multinomial(logits: lastLogit.squeezingShape(at: 1), numSamples: Tensor(1)) print(encoder.decode(tokens[0].makeNumpyArray()), terminator: "") } print()