├── .gitignore ├── .swiftformat ├── LICENSE ├── Package.swift ├── README.md ├── Sources ├── C │ ├── CBlas │ │ ├── cblas.h │ │ └── module.modulemap │ └── CLapack │ │ ├── lapacke.h │ │ └── module.modulemap ├── Debug │ └── main.swift └── NDArray │ ├── Extensions │ └── Expressible.swift │ ├── NDArray.swift │ ├── NDArrayTypes │ ├── BaseNDArray │ │ ├── ArrayShape.swift │ │ ├── BaseNDArray.swift │ │ └── Dimension.swift │ └── ScalarNDArray.swift │ ├── Ops │ ├── add.swift │ ├── broadcast.swift │ ├── description.swift │ ├── div.swift │ ├── dot.swift │ ├── elementWise.swift │ ├── max.swift │ ├── mean.swift │ ├── min.swift │ ├── mul.swift │ ├── reduce.swift │ ├── sub.swift │ ├── subscript.swift │ └── sum.swift │ ├── Protocols │ ├── Divisible.swift │ ├── NDArrayFloatingPoint.swift │ └── NDArrayProtocol.swift │ ├── Ref.swift │ └── utils.swift └── Tests ├── LinuxMain.swift └── NDArrayTests ├── DimensionTests.swift ├── NDArrayTests.swift └── XCTestManifests.swift /.gitignore: -------------------------------------------------------------------------------- 1 | # Xcode 2 | # 3 | # gitignore contributors: remember to update Global/Xcode.gitignore, Objective-C.gitignore & Swift.gitignore 4 | 5 | ## User settings 6 | xcuserdata/ 7 | 8 | ## compatibility with Xcode 8 and earlier (ignoring not required starting Xcode 9) 9 | *.xcscmblueprint 10 | *.xccheckout 11 | 12 | ## compatibility with Xcode 3 and earlier (ignoring not required starting Xcode 4) 13 | build/ 14 | DerivedData/ 15 | *.moved-aside 16 | *.pbxuser 17 | !default.pbxuser 18 | *.mode1v3 19 | !default.mode1v3 20 | *.mode2v3 21 | !default.mode2v3 22 | *.perspectivev3 23 | !default.perspectivev3 24 | 25 | ## Obj-C/Swift specific 26 | *.hmap 27 | 28 | ## App packaging 29 | *.ipa 30 | *.dSYM.zip 31 | *.dSYM 32 | 33 | ## Playgrounds 34 | timeline.xctimeline 35 | playground.xcworkspace 36 | 37 | # Swift Package Manager 38 | # 39 | # Add this line if you want to avoid checking in source code from Swift Package Manager dependencies. 40 | # Packages/ 41 | # Package.pins 42 | # Package.resolved 43 | # *.xcodeproj 44 | # 45 | # Xcode automatically generates this directory with a .xcworkspacedata file and xcuserdata 46 | # hence it is not needed unless you have added a package configuration file to your project 47 | # .swiftpm 48 | 49 | .build/ 50 | 51 | # CocoaPods 52 | # 53 | # We recommend against adding the Pods directory to your .gitignore. However 54 | # you should judge for yourself, the pros and cons are mentioned at: 55 | # https://guides.cocoapods.org/using/using-cocoapods.html#should-i-check-the-pods-directory-into-source-control 56 | # 57 | # Pods/ 58 | # 59 | # Add this line if you want to avoid checking in source code from the Xcode workspace 60 | # *.xcworkspace 61 | 62 | # Carthage 63 | # 64 | # Add this line if you want to avoid checking in source code from Carthage dependencies. 65 | # Carthage/Checkouts 66 | 67 | Carthage/Build/ 68 | 69 | # Accio dependency management 70 | Dependencies/ 71 | .accio/ 72 | 73 | # fastlane 74 | # 75 | # It is recommended to not store the screenshots in the git repo. 76 | # Instead, use fastlane to re-generate the screenshots whenever they are needed. 77 | # For more information about the recommended setup visit: 78 | # https://docs.fastlane.tools/best-practices/source-control/#source-control 79 | 80 | fastlane/report.xml 81 | fastlane/Preview.html 82 | fastlane/screenshots/**/*.png 83 | fastlane/test_output 84 | 85 | # Code Injection 86 | # 87 | # After new code Injection tools there's a generated folder /iOSInjectionProject 88 | # https://github.com/johnno1962/injectionforxcode 89 | 90 | iOSInjectionProject/ 91 | 92 | # editor 93 | .vscode -------------------------------------------------------------------------------- /.swiftformat: -------------------------------------------------------------------------------- 1 | --disable unusedArguments 2 | --disable spaceAroundOperators -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License 2 | 3 | Copyright (c) 2010-2019 Google, Inc. http://angularjs.org 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in 13 | all copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 21 | THE SOFTWARE. 22 | -------------------------------------------------------------------------------- /Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version:5.1 2 | // The swift-tools-version declares the minimum version of Swift required to build this package. 3 | 4 | import PackageDescription 5 | 6 | let package = Package( 7 | name: "NDArray", 8 | products: [ 9 | // Products define the executables and libraries produced by a package, and make them visible to other packages. 10 | .library( 11 | name: "NDArray", 12 | targets: ["NDArray"] 13 | ), 14 | .library( 15 | name: "Debug", 16 | targets: ["Debug"] 17 | ), 18 | ], 19 | dependencies: [ 20 | // Dependencies declare other packages that this package depends on. 21 | // .package(url: "../NDArray-clibs", from: "2.0.0"), 22 | ], 23 | targets: [ 24 | // Targets are the basic building blocks of a package. A target can define a module or a test suite. 25 | // Targets can depend on other targets in this package, and on products in packages which this package depends on. 26 | // .systemLibrary( 27 | // name: "CBlas", 28 | // path: "Sources/C/CBlas", 29 | // pkgConfig: "blas", 30 | // providers: [ 31 | // .apt(["gfortran", "liblapack3", "liblapacke", "liblapacke-dev", "libopenblas-base", "libopenblas-dev"]), 32 | // .brew(["homebrew/dupes/lapack", "homebrew/science/openblas"]), 33 | // ] 34 | // ), 35 | // .systemLibrary( 36 | // name: "CLapack", 37 | // path: "Sources/C/CLapack", 38 | // pkgConfig: "lapack", 39 | // providers: [ 40 | // .apt(["gfortran", "liblapack3", "liblapacke", "liblapacke-dev", "libopenblas-base", "libopenblas-dev"]), 41 | // .brew(["homebrew/dupes/lapack", "homebrew/science/openblas"]), 42 | // ] 43 | // ), 44 | .target( 45 | name: "NDArray", 46 | dependencies: [] 47 | // dependencies: ["CBlas", "CLapack"] 48 | ), 49 | .target( 50 | name: "Debug", 51 | dependencies: ["NDArray"] 52 | ), 53 | .testTarget( 54 | name: "NDArrayTests", 55 | dependencies: ["NDArray"] 56 | ), 57 | ] 58 | ) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # NDArray 2 | 3 | NDArray is a multidimensional array library written in Swift that aims to become the equivalent of `numpy` in Swift's emerging data science ecosystem. This project is in a very early stage and has a long but exciting road ahead! 4 | 5 | ## Goals 6 | 7 | 1. Have an efficient multidimensional array interface with common things like indexing, slicing, broadcasting, etc. 8 | 2. Make `NDArray` and its operations `differentiable` so its usable along with Swift for TensorFlow. 9 | 3. Create specialized implementations of linear algebra operations for NDArrays containing numeric types using BLAS, LAPACK, Accelerate, or [MLIR](https://docs.google.com/document/d/1UIPWl4lvBTozBD5OQ9SrxgcM7rA4pODMOjqQv3tm57w) depending on the environment. 10 | 11 | ## Tutorials ![](https://www.tensorflow.org/images/colab_logo_32px.png) 12 | 13 | Tutorial | Last Updated | 14 | -------- | ------------ | 15 | [Basic API](https://colab.research.google.com/drive/1aULWtrtj6WsNeJe_vBnr_hswy0JIYDt_) | August 13 2019 | 16 | 17 | 18 | ## Installation 19 | You can install it using SwiftPM: 20 | ```swift 21 | .package(url: "https://github.com/cgarciae/NDArray", from: "0.0.20") 22 | ``` 23 | It might work on other compatible package managers. This package is only tested in Swift 5.1, compatibility with previous version is not guaranteed. 24 | 25 | ## Example 26 | `NDArray` is a generic container type just like `Array` with the difference that its multidimensional. If its elements conform to certain protocols then certain methods and operators like `+`, `-`, `*`, etc, can be used to efficiently perform computations of the whole collection. 27 | ```swift 28 | import NDArray 29 | 30 | let a = NDArray([ 31 | [1, 2, 3], 32 | [4, 5, 6], 33 | ]) 34 | let b = NDArray([ 35 | [7, 8, 9], 36 | [10, 11, 12], 37 | ]) 38 | 39 | print((a + b) * a) 40 | /* 41 | NDArray[2, 3]([ 42 | [8, 20, 36], 43 | [56, 80, 108], 44 | ]) 45 | */ 46 | ``` 47 | Here we see that the outcome of `(a + b) * a` is also and `NDArray` of `Int` with shape `[2, 3]`. To use operators like `+` and `*` with NDArrays containing your custom types you just have to make them conform to the proper protocols. For example: 48 | ```swift 49 | import NDArray 50 | 51 | struct Point: AdditiveArithmetic { 52 | let x: Float 53 | let y: Float 54 | ... 55 | } 56 | 57 | let a = NDArray([Point(x: 1, y: 2), Point(x: 2, y: 3)]) 58 | let b = NDArray([Point(x: 4, y: 5), Point(x: 6, y: 7)]) 59 | 60 | print(a + b) 61 | /* 62 | NDArray[2]([Point(x: 5.0, y: 7.0), Point(x: 8.0, y: 10.0)]) 63 | */ 64 | ``` 65 | You can also apply generic transformations over the data, the previous could have been written as: 66 | ```swift 67 | elementwise(a, b, apply: +) 68 | // or 69 | elementwise(a, b) { $0 + $1 } 70 | ``` 71 | For heavy computation you can use the parallelized version: 72 | ```swift 73 | elementwiseInParallel(a, b) { 74 | // code 75 | return c 76 | } 77 | ``` 78 | In the future `NDArray` should be able to estimate the best strategy (serial/parallelized) based on the type and size of the data. 79 | 80 | ## Goals Discussion 81 | Except for the Basic API, NDArray's Automatic Differentiation and Linear Algebra Optimization capabilities should be opt-in so all users can have access to the library regardless of their environment, i.e. iOS developers should be able to use it even if they don't have access to TensorFlow's compiler or the Lineal Algebra infrastructure. 82 | 83 | #### Basic API 84 | The first goal is the definition of the library's basic API using pure Swift with no extra optimization or differentiable capabilities. iOS/OSX developers should be able to use the basic API without additional setup. It will also be important to keeping the NDArray's API in close coordination with Swift for TensorFlow's Tensor API to promote knowledge reuse and free documentation if possible. 85 | 86 | #### Automatic Differentiation 87 | The second goal is an obvious must have, Swift for TensorFlow's compiler with automatic differentiation is arguably the future of ML and we should use it. 88 | 89 | #### Linear Algebra Optimization 90 | The third goal is what you would expect from any HPC numeric library, the strategy would be to specialize functions/operations for numeric types by using BLAS, LAPACK, Accelerate, or MLIR to speed computation. On the other hand, if successfully integrated with [MLIR](https://docs.google.com/document/d/1UIPWl4lvBTozBD5OQ9SrxgcM7rA4pODMOjqQv3tm57w), BLAS and LAPACK might not be necessary and NDArray could easily become one of the most performant numeric libraries out there. 91 | 92 | 93 | ## Roadmap 94 | ##### 0.1: Basic API 95 | - [x] Indexing 96 | - [x] Dimension Slicing 97 | - [x] Dimension Filtering by Indexes 98 | - [x] Dimension Masking 99 | - [x] SqueezeAxis 100 | - [x] NewAxis 101 | - [x] Assignment 102 | - [x] Broadcasting 103 | - [x] Pretty Print 104 | - [x] Elementwise Operations 105 | - [x] Basic Operators: `+`, `-`, `*`, `\` 106 | - [x] Reduction Operations (reduce, sum, mean, max, min) 107 | - [ ] Concatenation Operations (concat, stack, hstack, vstack) 108 | - [ ] Subscript Bound Checks 109 | - [ ] Fancy Indexing 110 | - [ ] > 95% Coverage 111 | - [ ] Documentation 112 | ##### 0.2: Differentiable Programming 113 | This can actually be started at any point, although it wont be that useful until various operations like `dot` or reductions like `sum` or `mean` are implemented. 114 | - [ ] Conform `NDArray` to `Differentiable` 115 | - [ ] Make `NDArrays` operations differentiable. 116 | ##### 0.3: Linear Algebra Optimization 117 | - [x] Link BLAS and LAPACK 118 | - [ ] Specialize operations using BLAS, LAPACK, Accelerate, or [MLIR](https://docs.google.com/document/d/1UIPWl4lvBTozBD5OQ9SrxgcM7rA4pODMOjqQv3tm57w) 119 | - [ ] `dot` 120 | - [ ] ... 121 | 122 | ## Meta 123 | Cristian Garcia – cgarcia.e88@gmail.com 124 | 125 | Distributed under the MIT license. See LICENSE for more information. 126 | 127 | -------------------------------------------------------------------------------- /Sources/C/CBlas/cblas.h: -------------------------------------------------------------------------------- 1 | #include -------------------------------------------------------------------------------- /Sources/C/CBlas/module.modulemap: -------------------------------------------------------------------------------- 1 | module CBlas { 2 | umbrella header "cblas.h" 3 | link "blas" 4 | } 5 | 6 | -------------------------------------------------------------------------------- /Sources/C/CLapack/lapacke.h: -------------------------------------------------------------------------------- 1 | #include -------------------------------------------------------------------------------- /Sources/C/CLapack/module.modulemap: -------------------------------------------------------------------------------- 1 | module CLapack { 2 | umbrella header "lapacke.h" 3 | link "lapack" 4 | } 5 | 6 | -------------------------------------------------------------------------------- /Sources/Debug/main.swift: -------------------------------------------------------------------------------- 1 | import NDArray 2 | 3 | let a = NDArray([ 4 | [1, 2], 5 | [5, 6], 6 | ]) 7 | 8 | let b = NDArray([ 9 | [1, 10], 10 | [100, 1000], 11 | ]) 12 | 13 | let c = elementwise(a, b) { $0 + $1 } 14 | 15 | print(c) 16 | 17 | // let a = NDArray([ 18 | // [1, 2, 3], 19 | // [4, 5, 6], 20 | // ]).transposed([1, 0]).copy() 21 | 22 | // let s = indexSequence(range: 3 ..< 10, shape: [10, 2, 5]) 23 | 24 | // for x in s { 25 | // print(x.rectangularIndex) 26 | // } 27 | 28 | // print(NDArray(0)) 29 | // print(NDArray([1, 2, 3, 4])) 30 | // print(NDArray([ 31 | // [1, 2, 3], 32 | // [4, 5, 6], 33 | // ])) 34 | // print(NDArray([ 35 | // [1, 2, 3], 36 | // [4, 5, 6], 37 | // ]).transposed([1, 0])) 38 | // print(NDArray( 39 | // [ 40 | // [ 41 | // [1, 2, 3], 42 | // [4, 5, 6], 43 | // ], 44 | // [ 45 | // [7, 8, 9], 46 | // [10, 11, 12], 47 | // ], 48 | // ] 49 | // )) 50 | 51 | // print(NDArray([ 52 | // [1, 2, 3, 4, 5, 6, 7], 53 | // [4, 5, 6, 7, 8, 9, 10], 54 | // ])[0..., 1 ..< 5][0..., ..<3]) 55 | 56 | // print(NDArray([ 57 | // [1, 2, 3, 4, 5, 6, 7], 58 | // [4, 5, 6, 7, 8, 9, 10], 59 | // ])[0..., ((-1)...).stride(2)]) -------------------------------------------------------------------------------- /Sources/NDArray/Extensions/Expressible.swift: -------------------------------------------------------------------------------- 1 | // This strategy was taken from TensorFlow 2 | // See: https://github.com/tensorflow/swift-apis/blob/728a2d6c2146d12d8a16428f6c9a3bce24492bb7/Sources/TensorFlow/Core/Tensor.swift#L278 3 | 4 | // At first is doesn't seem ideal since it would be nice to enable ExpressibleBy(Int|Float|Bool)Literal to 5 | // be able to write things like x[1, 0...] = 5, however, while this could work in the short-term when NDArray 6 | // is just a struct, it might not work in the future when it might turn into a protocol (to house SparseArray, DenseArray, etc) 7 | 8 | @frozen 9 | public struct _NDArrayElementLiteral { 10 | @usableFromInline let array: [Scalar] 11 | @usableFromInline let shape: [Int] 12 | } 13 | 14 | extension _NDArrayElementLiteral: ExpressibleByBooleanLiteral 15 | where Scalar: ExpressibleByBooleanLiteral { 16 | public typealias BooleanLiteralType = Scalar.BooleanLiteralType 17 | @inlinable 18 | public init(booleanLiteral: BooleanLiteralType) { 19 | array = [Scalar(booleanLiteral: booleanLiteral)] 20 | shape = [] 21 | } 22 | } 23 | 24 | extension _NDArrayElementLiteral: ExpressibleByIntegerLiteral 25 | where Scalar: ExpressibleByIntegerLiteral { 26 | public typealias IntegerLiteralType = Scalar.IntegerLiteralType 27 | @inlinable 28 | public init(integerLiteral: IntegerLiteralType) { 29 | array = [Scalar(integerLiteral: integerLiteral)] 30 | shape = [] 31 | } 32 | } 33 | 34 | extension _NDArrayElementLiteral: ExpressibleByFloatLiteral 35 | where Scalar: ExpressibleByFloatLiteral { 36 | public typealias FloatLiteralType = Scalar.FloatLiteralType 37 | 38 | @inlinable 39 | public init(floatLiteral: FloatLiteralType) { 40 | array = [Scalar(floatLiteral: floatLiteral)] 41 | shape = [] 42 | } 43 | } 44 | 45 | extension _NDArrayElementLiteral: ExpressibleByArrayLiteral { 46 | public typealias ArrayLiteralElement = _NDArrayElementLiteral 47 | @inlinable 48 | public init(arrayLiteral elements: _NDArrayElementLiteral...) { 49 | array = elements.reduce([Scalar]()) { array, element in 50 | array + element.array 51 | } 52 | shape = [elements.count] + elements[0].shape 53 | } 54 | } 55 | 56 | extension NDArray: ExpressibleByArrayLiteral { 57 | /// The type of the elements of an array literal. 58 | public typealias ArrayLiteralElement = _NDArrayElementLiteral 59 | 60 | /// Creates a ndarray initialized with the given elements. 61 | /// - Note: This is for conversion from ndarray element literals. This is a 62 | /// separate method because `ShapedArray` initializers need to call it. 63 | @inlinable 64 | internal init(_tensorElementLiterals elements: [_NDArrayElementLiteral]) { 65 | let array = elements.reduce([Scalar]()) { array, element in 66 | array + element.array 67 | } 68 | let shape = [elements.count] + elements[0].shape 69 | 70 | self = NDArray(array, shape: shape) 71 | } 72 | 73 | /// Creates a ndarray initialized with the given elements. 74 | @inlinable 75 | public init(arrayLiteral elements: _NDArrayElementLiteral...) { 76 | precondition(!elements.isEmpty, "Cannot create a 'NDArray' with no elements.") 77 | self.init(_tensorElementLiterals: elements) 78 | } 79 | } -------------------------------------------------------------------------------- /Sources/NDArray/NDArray.swift: -------------------------------------------------------------------------------- 1 | 2 | 3 | public struct NDArray { 4 | public let shape: [Int] 5 | 6 | public typealias ScalarGetter = ((Int, UnsafeMutableBufferPointer)) -> Scalar 7 | public typealias ScalarSetter = ((Int, UnsafeMutableBufferPointer), Scalar) -> Void 8 | public typealias ScalarGetterSetter = ((Int, UnsafeMutableBufferPointer), (Scalar) -> Scalar) -> Void 9 | 10 | public var anyNDArray: AnyNDArray 11 | 12 | public init(_ ndarray: N) where N.Scalar == Scalar { 13 | shape = ndarray.shape 14 | anyNDArray = AnyNDArray(ndarray) 15 | } 16 | 17 | internal init(_ anyNDArray: AnyNDArray) { 18 | shape = anyNDArray.shape 19 | self.anyNDArray = anyNDArray 20 | } 21 | 22 | init(_ ndarray: NDArray) { 23 | self = ndarray 24 | } 25 | 26 | public init(_ data: [Any], shape: [Int]? = nil) { 27 | let (flatData, calculatedShape): ([Scalar], [Int]) = flattenArrays(data) 28 | 29 | precondition( 30 | calculatedShape.product() == flatData.count, 31 | "All sub-arrays in data must have equal length. Calculated shape: \(calculatedShape), \(flatData)" 32 | ) 33 | 34 | if let shape = shape { 35 | precondition( 36 | shape.product() == flatData.count, 37 | "Invalid shape, number of elements" 38 | ) 39 | } 40 | 41 | let data = Ref(flatData) 42 | let arrayShape = ArrayShape(shape ?? calculatedShape) 43 | 44 | self = NDArray(BaseNDArray(data, shape: arrayShape)) 45 | } 46 | 47 | public init(_ data: Scalar) { 48 | self = NDArray(ScalarNDArray( 49 | data, 50 | shape: [] 51 | )) 52 | } 53 | 54 | public func subscript_get(_ ranges: [ArrayRange]) -> NDArray { 55 | anyNDArray.subscript_get(ranges) 56 | } 57 | 58 | public mutating func subscript_set(_ ranges: [ArrayRange], _ value: NDArray) -> NDArray { 59 | if !isKnownUniquelyReferenced(&anyNDArray) { 60 | self = anyNDArray.copy() 61 | } 62 | 63 | self = anyNDArray.subscript_set(ranges, value) 64 | 65 | return self 66 | } 67 | 68 | public func linearIndex(at indexes: [Int]) -> Int { 69 | anyNDArray.linearIndex(indexes) 70 | } 71 | 72 | public func dataValue(at indexes: [Int]) -> Scalar { 73 | anyNDArray.dataValue(indexes) 74 | } 75 | 76 | public func withScalarGetter(_ body: (@escaping ScalarGetter) -> Void) { 77 | anyNDArray.withScalarGetter(body) 78 | } 79 | 80 | public mutating func withScalarSetter(_ body: (@escaping ScalarSetter) -> Void) { 81 | anyNDArray.withScalarSetter(body) 82 | } 83 | 84 | public mutating func withScalarGetterSetter(_ body: (@escaping ScalarGetterSetter) -> Void) { 85 | anyNDArray.withScalarGetterSetter(body) 86 | } 87 | 88 | public func transposed(_ permutations: [Int]) -> NDArray { 89 | anyNDArray.transposed(permutations) 90 | } 91 | 92 | public func tiled(by amounts: [Int]) -> NDArray { 93 | anyNDArray.tiled(amounts) 94 | } 95 | 96 | public func expandDimensions(axis: Int) -> NDArray { 97 | anyNDArray.expandDimensions(axis) 98 | } 99 | 100 | public func scalarized() -> Scalar { 101 | anyNDArray.scalarized() 102 | } 103 | 104 | public func copy() -> NDArray { 105 | NDArray(anyNDArray.copy()) 106 | } 107 | 108 | public func baseCopy() -> BaseNDArray { 109 | anyNDArray.baseCopy() 110 | } 111 | 112 | public func toArray(_: T.Type) -> T { 113 | let cp: BaseNDArray = baseCopy() 114 | var array: [Any] = cp.data.value 115 | 116 | for n in cp.shape.reversed().dropLast() { 117 | array = array.chunked(into: n) as [Any] 118 | } 119 | 120 | return array as! T 121 | } 122 | } 123 | 124 | public class AnyNDArray { 125 | public let shape: [Int] 126 | 127 | let linearIndex: ([Int]) -> Int 128 | let dataValue: ([Int]) -> Scalar 129 | let subscript_get: ([ArrayRange]) -> NDArray 130 | let subscript_set: ([ArrayRange], NDArray) -> NDArray 131 | let withScalarGetter: ((@escaping NDArray.ScalarGetter) -> Void) -> Void 132 | let withScalarSetter: ((@escaping NDArray.ScalarSetter) -> Void) -> Void 133 | let withScalarGetterSetter: ((@escaping NDArray.ScalarGetterSetter) -> Void) -> Void 134 | let transposed: ([Int]) -> NDArray 135 | let tiled: ([Int]) -> NDArray 136 | let expandDimensions: (Int) -> NDArray 137 | let scalarized: () -> Scalar 138 | let copy: () -> NDArray 139 | let baseCopy: () -> BaseNDArray 140 | 141 | // 142 | @usableFromInline 143 | let isSetable: () -> Bool 144 | 145 | public init(_ ndarray: N) where N.Scalar == Scalar { 146 | var ndarray = ndarray 147 | 148 | shape = ndarray.shape 149 | linearIndex = { ndarray.linearIndex(at: $0) } 150 | dataValue = { ndarray.dataValue(at: $0) } 151 | subscript_get = { ndarray.subscript_get($0) } 152 | subscript_set = { 153 | ndarray.subscript_set($0, $1) 154 | } 155 | withScalarGetter = { ndarray.withScalarGetter($0) } 156 | withScalarSetter = { ndarray.withScalarSetter($0) } 157 | withScalarGetterSetter = { ndarray.withScalarGetterSetter($0) } 158 | transposed = { ndarray.transposed($0) } 159 | tiled = { ndarray.tiled(by: $0) } 160 | expandDimensions = { ndarray.expandDimensions(axis: $0) } 161 | scalarized = { ndarray.scalarized() } 162 | copy = { ndarray.copy() } 163 | baseCopy = { ndarray.baseCopy() } 164 | isSetable = { ndarray is SetableNDArray } 165 | } 166 | } -------------------------------------------------------------------------------- /Sources/NDArray/NDArrayTypes/BaseNDArray/ArrayShape.swift: -------------------------------------------------------------------------------- 1 | 2 | 3 | public struct ArrayShape { 4 | @usableFromInline let dimensions: [DimensionProtocol] 5 | @usableFromInline let linearMemoryOffset: Int 6 | @usableFromInline let dimensionLengths: [Int] 7 | @usableFromInline let dimensionStrides: [Int] 8 | @usableFromInline let isOriginalShape: Bool 9 | 10 | @inlinable 11 | public init(_ shape: [Int]) { 12 | let dimensionStrides = getDimensionStrides(of: shape) 13 | self.dimensionStrides = dimensionStrides 14 | 15 | dimensionLengths = shape 16 | linearMemoryOffset = 0 17 | 18 | dimensions = (0 ..< shape.count) 19 | .map { i -> DimensionProtocol in 20 | 21 | if shape[i] == 1 { 22 | return SingularDimension() 23 | } else { 24 | return Dimension(length: shape[i], memory_stride: dimensionStrides[i]) 25 | } 26 | } 27 | 28 | isOriginalShape = true 29 | } 30 | 31 | @usableFromInline init(_ dimensions: [DimensionProtocol], linearMemoryOffset: Int) { 32 | self.dimensions = dimensions 33 | self.linearMemoryOffset = linearMemoryOffset 34 | dimensionLengths = dimensions.map { $0.length } 35 | dimensionStrides = getDimensionStrides(of: dimensionLengths) 36 | isOriginalShape = linearMemoryOffset == 0 && 37 | dimensions.lazy.map { $0 is UnmodifiedDimension }.all() 38 | } 39 | 40 | @inlinable 41 | public func linearIndex(of indexes: UnsafeMutableBufferPointer) -> Int { 42 | let partialIndex = zip(indexes, dimensions) 43 | .lazy 44 | .map { index, dimension in 45 | dimension.strideValue(of: index) 46 | } 47 | .sum() 48 | 49 | return partialIndex + linearMemoryOffset 50 | } 51 | } -------------------------------------------------------------------------------- /Sources/NDArray/NDArrayTypes/BaseNDArray/BaseNDArray.swift: -------------------------------------------------------------------------------- 1 | 2 | 3 | public struct BaseNDArray: NDArrayProtocol, SetableNDArray { 4 | public var data: Ref<[Scalar]> 5 | public var arrayShape: ArrayShape 6 | 7 | public var shape: [Int] { arrayShape.dimensionLengths } 8 | 9 | @inlinable 10 | public init(_ data: Ref<[Scalar]>, shape: ArrayShape) { 11 | self.data = data 12 | arrayShape = shape 13 | } 14 | 15 | // public init(_ data: [Any], shape: [Int]? = nil) { 16 | // let (flatData, calculatedShape): ([Scalar], [Int]) = flattenArrays(data) 17 | 18 | // precondition( 19 | // calculatedShape.product() == flatData.count, 20 | // "All sub-arrays in data must have equal length. Calculated shape: \(calculatedShape), \(flatData)" 21 | // ) 22 | 23 | // if let shape = shape { 24 | // precondition( 25 | // shape.product() == flatData.count, 26 | // "Invalid shape, number of elements" 27 | // ) 28 | // } 29 | 30 | // arrayShape = ArrayShape(shape ?? calculatedShape) 31 | // self.data = Ref(flatData) 32 | // } 33 | 34 | // @usableFromInline 35 | // internal init(_ data: [Scalar], shape: ArrayShape) { 36 | // arrayShape = shape 37 | // self.data = Ref(data) 38 | // } 39 | 40 | // public init(_ data: Scalar) { 41 | // arrayShape = ArrayShape([DimensionProtocol](), linearMemoryOffset: 0) 42 | // self.data = Ref([data]) 43 | // } 44 | 45 | @inlinable 46 | public func linearIndex(at indexes: UnsafeMutableBufferPointer) -> Int { 47 | arrayShape.linearIndex(of: indexes) 48 | } 49 | 50 | @inlinable 51 | public func linearIndex(at indexes: [Int]) -> Int { 52 | var indexes = indexes 53 | 54 | return indexes.withUnsafeMutableBufferPointer { indexes in 55 | linearIndex(at: indexes) 56 | } 57 | } 58 | 59 | @inlinable 60 | public func dataValue(at indexes: UnsafeMutableBufferPointer) -> Scalar { 61 | return data.value[ 62 | arrayShape.linearIndex(of: indexes) 63 | ] 64 | } 65 | 66 | @inlinable 67 | public func dataValue(at indexes: [Int]) -> Scalar { 68 | var indexes = indexes 69 | 70 | return indexes.withUnsafeMutableBufferPointer { indexes in 71 | dataValue(at: indexes) 72 | } 73 | } 74 | 75 | @inlinable 76 | public func withScalarGetter(_ body: (@escaping NDArray.ScalarGetter) -> Void) { 77 | data.value.withUnsafeBufferPointer { data in 78 | 79 | body { _, rectIndex in 80 | let index = self.arrayShape.linearIndex(of: rectIndex) 81 | return data[index] 82 | } 83 | } 84 | } 85 | 86 | @inlinable 87 | public mutating func withScalarSetter(_ body: (@escaping NDArray.ScalarSetter) -> Void) { 88 | data.value.withUnsafeMutableBufferPointer { dataIn in 89 | var data = dataIn 90 | defer { dataIn = data } 91 | 92 | body { [self] indexer, value in 93 | let (_, rectangularIndex) = indexer 94 | let index = self.arrayShape.linearIndex(of: rectangularIndex) 95 | 96 | data[index] = value 97 | } 98 | } 99 | } 100 | 101 | @inlinable 102 | public mutating func withScalarGetterSetter(_ body: (@escaping NDArray.ScalarGetterSetter) -> Void) { 103 | data.value.withUnsafeMutableBufferPointer { dataIn in 104 | var data = dataIn 105 | defer { dataIn = data } 106 | 107 | body { [self] index, f in 108 | let index = self.arrayShape.linearIndex(of: index.1) 109 | let value = f(data[index]) 110 | 111 | data[index] = value 112 | } 113 | } 114 | } 115 | 116 | @inlinable 117 | public func subscript_get(_ ranges: [ArrayRange]) -> NDArray { 118 | var ranges = ranges 119 | 120 | let nEllipsis = ranges.filter(isEllipsis).count 121 | 122 | precondition(nEllipsis <= 1, "A maximum of 1 .ellipsis can be used, got \(ranges)") 123 | 124 | if nEllipsis == 1 { 125 | let ellipsisIndex = ranges.firstIndex(where: isEllipsis)! 126 | let nAll = 1 + shape.count - ranges.count 127 | 128 | ranges.remove(at: ellipsisIndex) 129 | 130 | for _ in 0 ..< nAll { 131 | ranges.insert(.all, at: ellipsisIndex) 132 | } 133 | } 134 | 135 | precondition(shape.count >= ranges.count) 136 | 137 | var dimensions = arrayShape.dimensions 138 | var linearMemoryOffset = arrayShape.linearMemoryOffset 139 | var dimensionToBeRemoved = [Int]() 140 | var dimensionToBeAdded = [Int: DimensionProtocol]() 141 | 142 | for (i, range) in ranges.enumerated() { 143 | switch range { 144 | case let .index(index): 145 | let index = index < 0 ? dimensions[i].length + index : index 146 | 147 | linearMemoryOffset += dimensions[i].strideValue(of: index) 148 | dimensionToBeRemoved.append(i) 149 | 150 | case let .slice(start: start, end: end, stride: stride): 151 | 152 | if start == 0, end == nil || end! == dimensions[i].length, stride == 1 { 153 | continue 154 | } 155 | 156 | dimensions[i] = dimensions[i].sliced( 157 | start: start, 158 | end: end, 159 | stride: stride 160 | ) 161 | case let .filter(indexes): 162 | dimensions[i] = dimensions[i].select(indexes: indexes) 163 | 164 | case .all: 165 | continue 166 | case .squeezeAxis: 167 | precondition( 168 | dimensions[i].length == 1, 169 | "Cannot squeeze dimension \(i) of \(shape), expected 1 got \(shape[i])" 170 | ) 171 | 172 | linearMemoryOffset += dimensions[i].strideValue(of: 0) 173 | dimensionToBeRemoved.append(i) 174 | 175 | case .newAxis: 176 | dimensionToBeAdded[i] = SingularDimension() 177 | 178 | case .ellipsis: 179 | fatalError("Ellipsis should be expand as a series of .all expressions") 180 | } 181 | } 182 | 183 | // TODO: this implementation is not correct due the fact the the length of dimension is changing 184 | // A correct way to implement this would be to do the operations sorted by the index 185 | // from high to low. 186 | dimensions = dimensions 187 | .enumerated() 188 | .filter { i, d in !dimensionToBeRemoved.contains(i) } 189 | .map { i, d in d } 190 | 191 | for (i, dimension) in dimensionToBeAdded { 192 | dimensions.insert(dimension, at: i) 193 | } 194 | 195 | let ndarray = BaseNDArray( 196 | data, 197 | shape: ArrayShape( 198 | dimensions, 199 | linearMemoryOffset: linearMemoryOffset 200 | ) 201 | ) 202 | 203 | if ndarray.shape.isEmpty { 204 | return NDArray(ScalarNDArray( 205 | ndarray.scalarized(), 206 | shape: [] 207 | )) 208 | } else { 209 | return NDArray(ndarray) 210 | } 211 | } 212 | 213 | public mutating func subscript_set(_ ranges: [ArrayRange], _ ndarray: NDArray) -> NDArray { 214 | var ndarray = ndarray 215 | 216 | if !isKnownUniquelyReferenced(&data) { 217 | self = baseCopy() 218 | } 219 | 220 | var ndarrayView = subscript_get(ranges) 221 | let nElements = ndarrayView.shape.product() 222 | 223 | if ndarrayView.shape != ndarray.shape { 224 | (ndarrayView, ndarray) = broadcast(ndarrayView, and: ndarray) 225 | } 226 | 227 | let ndarrayViewShape = ndarrayView.shape 228 | 229 | ndarrayView.withScalarSetter { viewScalarSetter in 230 | ndarray.withScalarGetter { ndarrayScalarGetter in 231 | for index in indexSequence(range: 0 ..< nElements, shape: ndarrayViewShape) { 232 | let value = ndarrayScalarGetter(index) 233 | viewScalarSetter(index, value) 234 | } 235 | } 236 | } 237 | 238 | return NDArray(self) 239 | } 240 | 241 | public func tiled(by repetitions: [Int]) -> NDArray { 242 | var dimensions = arrayShape.dimensions 243 | 244 | for i in 0 ..< dimensions.count { 245 | if repetitions[i] > 1 { 246 | dimensions[i] = dimensions[i].tiled(repetitions[i]) 247 | } 248 | } 249 | 250 | return NDArray(BaseNDArray( 251 | data, 252 | shape: ArrayShape( 253 | dimensions, 254 | linearMemoryOffset: arrayShape.linearMemoryOffset 255 | ) 256 | )) 257 | } 258 | 259 | public func expandDimensions(axis: Int) -> NDArray { 260 | var dimensions = arrayShape.dimensions 261 | 262 | dimensions.insert(SingularDimension(), at: axis) 263 | 264 | return NDArray(BaseNDArray( 265 | data, 266 | shape: ArrayShape( 267 | dimensions, 268 | linearMemoryOffset: arrayShape.linearMemoryOffset 269 | ) 270 | )) 271 | } 272 | 273 | public func scalarized() -> Scalar { 274 | precondition(shape == [] || shape == [1], "Cannot convert non-scalar NDArray to scalar, got shape \(shape)") 275 | 276 | return dataValue(at: []) 277 | } 278 | 279 | public func copy() -> NDArray { 280 | let nElements = shape.product() 281 | 282 | let arrayC = [Scalar](unsafeUninitializedCapacity: nElements) { arrayC, count in 283 | count = nElements 284 | 285 | self.withScalarGetter { arrayAScalarGetter in 286 | for index in indexSequence(range: 0 ..< nElements, shape: shape) { 287 | arrayC[index.linearIndex] = arrayAScalarGetter(index) 288 | } 289 | } 290 | } 291 | 292 | return NDArray(BaseNDArray( 293 | Ref(arrayC), 294 | shape: ArrayShape(shape) 295 | )) 296 | } 297 | 298 | public func transposed(_ indexes: [Int]) -> NDArray { 299 | precondition(shape.count >= indexes.count) 300 | 301 | return NDArray(BaseNDArray( 302 | data, 303 | shape: ArrayShape( 304 | indexes.map { i in arrayShape.dimensions[i] }, 305 | linearMemoryOffset: arrayShape.linearMemoryOffset 306 | ) 307 | )) 308 | } 309 | } 310 | 311 | extension NDArrayProtocol { 312 | internal func baseCopy() -> BaseNDArray { 313 | let nElements = shape.product() 314 | 315 | let arrayC = [Scalar](unsafeUninitializedCapacity: nElements) { arrayC, count in 316 | count = nElements 317 | 318 | self.withScalarGetter { arrayAScalarGetter in 319 | for index in indexSequence(range: 0 ..< nElements, shape: shape) { 320 | arrayC[index.linearIndex] = arrayAScalarGetter(index) 321 | } 322 | } 323 | } 324 | 325 | return BaseNDArray( 326 | Ref(arrayC), 327 | shape: ArrayShape(shape) 328 | ) 329 | } 330 | } -------------------------------------------------------------------------------- /Sources/NDArray/NDArrayTypes/BaseNDArray/Dimension.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | 3 | public struct MemoryLayout { 4 | public let stride: Int 5 | public let length: Int 6 | 7 | fileprivate init(length: Int, stride: Int) { 8 | self.stride = stride 9 | self.length = length 10 | } 11 | } 12 | 13 | public protocol DimensionProtocol { 14 | var length: Int { get } 15 | var memory_layout: MemoryLayout { get } 16 | 17 | @inlinable 18 | func linearIndex(of: Int) -> Int 19 | } 20 | 21 | // public protocol SqueezedDimension: DimensionProtocol {} 22 | public protocol UnmodifiedDimension: DimensionProtocol {} 23 | 24 | extension DimensionProtocol { 25 | @inlinable 26 | public func strideValue(of index: Int) -> Int { 27 | linearIndex(of: index) * memory_layout.stride 28 | } 29 | } 30 | 31 | public struct Dimension: DimensionProtocol, UnmodifiedDimension { 32 | public let length: Int 33 | public let memory_layout: MemoryLayout 34 | 35 | public init(length: Int, memory_stride: Int) { 36 | self.length = length 37 | memory_layout = MemoryLayout(length: length, stride: memory_stride) 38 | } 39 | 40 | @inlinable 41 | public func linearIndex(of index: Int) -> Int { index } 42 | } 43 | 44 | public struct SingularDimension: DimensionProtocol, UnmodifiedDimension { 45 | public let length: Int = 1 46 | public let memory_layout: MemoryLayout 47 | 48 | public init() { 49 | memory_layout = MemoryLayout(length: 1, stride: 0) 50 | } 51 | 52 | @inlinable 53 | public func linearIndex(of index: Int) -> Int { 0 } 54 | } 55 | 56 | public struct SlicedDimension: DimensionProtocol { 57 | public let base: DimensionProtocol 58 | public let length: Int 59 | 60 | public let stride: Int 61 | public let start: Int 62 | public let end: Int 63 | 64 | public var memory_layout: MemoryLayout 65 | 66 | fileprivate init(base: DimensionProtocol, start: Int, end: Int, stride: Int) { 67 | self.base = base 68 | self.stride = stride 69 | self.start = start 70 | self.end = end 71 | 72 | if abs(end - start) % abs(stride) != 0 { 73 | length = (1 + abs(end - start) / abs(stride)) 74 | } else { 75 | length = (abs(end - start) / abs(stride)) 76 | } 77 | 78 | memory_layout = base.memory_layout 79 | } 80 | 81 | @inlinable 82 | public func linearIndex(of index: Int) -> Int { 83 | return base.linearIndex(of: start + index * stride) 84 | } 85 | } 86 | 87 | public struct InvertedDimension: DimensionProtocol { 88 | public let base: DimensionProtocol 89 | public let length: Int 90 | 91 | public var memory_layout: MemoryLayout 92 | 93 | fileprivate init(base: DimensionProtocol) { 94 | self.base = base 95 | 96 | length = base.length 97 | memory_layout = base.memory_layout 98 | } 99 | 100 | @inlinable 101 | public func linearIndex(of index: Int) -> Int { 102 | base.linearIndex(of: length - 1 - index) 103 | } 104 | } 105 | 106 | public struct TiledDimension: DimensionProtocol { 107 | public let base: DimensionProtocol 108 | public var length: Int 109 | public var repetitions: Int 110 | public var memory_layout: MemoryLayout 111 | 112 | fileprivate init(base: DimensionProtocol, repetitions: Int) { 113 | self.base = base 114 | self.repetitions = repetitions 115 | 116 | length = base.length * repetitions 117 | memory_layout = base.memory_layout 118 | } 119 | 120 | @inlinable 121 | public func linearIndex(of index: Int) -> Int { 122 | base.linearIndex(of: index % base.length) 123 | } 124 | } 125 | 126 | public struct FilteredDimension: DimensionProtocol { 127 | public let base: DimensionProtocol 128 | public var length: Int 129 | public var memory_layout: MemoryLayout 130 | public var indexes: [Int] 131 | 132 | fileprivate init(base: DimensionProtocol, indexes: [Int]) { 133 | self.base = base 134 | self.indexes = indexes.map { index in 135 | index < 0 ? index + base.length : index 136 | } 137 | 138 | length = indexes.count 139 | memory_layout = base.memory_layout 140 | } 141 | 142 | @inlinable 143 | public func linearIndex(of index: Int) -> Int { 144 | base.linearIndex(of: indexes[index]) 145 | } 146 | } 147 | 148 | extension DimensionProtocol { 149 | public func sliced(start: Int? = nil, end: Int? = nil, stride: Int = 1) -> DimensionProtocol { 150 | var start = start ?? (stride > 0 ? 0 : -1) 151 | var end = end ?? (stride > 0 ? length : 0) 152 | 153 | if start < 0 { 154 | start += length 155 | } 156 | if end < 0 { 157 | end += length 158 | } 159 | 160 | if stride < 0 { 161 | end -= 1 162 | } 163 | 164 | return SlicedDimension( 165 | base: self, 166 | start: start, 167 | end: end, 168 | stride: stride 169 | ) 170 | } 171 | 172 | public func tiled(_ repetitions: Int) -> DimensionProtocol { 173 | TiledDimension(base: self, repetitions: repetitions) 174 | } 175 | 176 | public func select(indexes: [Int]) -> DimensionProtocol { 177 | FilteredDimension(base: self, indexes: indexes) 178 | } 179 | } -------------------------------------------------------------------------------- /Sources/NDArray/NDArrayTypes/ScalarNDArray.swift: -------------------------------------------------------------------------------- 1 | public struct ScalarNDArray: NDArrayProtocol { 2 | public var data: Scalar 3 | public let shape: [Int] 4 | 5 | @inlinable 6 | public init(_ data: Scalar, shape: [Int]) { 7 | self.data = data 8 | self.shape = shape 9 | } 10 | 11 | public func subscript_get(_ ranges: [ArrayRange]) -> NDArray { 12 | var shape = self.shape 13 | var dimensionToBeRemoved = [Int]() 14 | var dimensionToBeAdded = [Int: Int]() 15 | 16 | for (i, range) in ranges.enumerated() { 17 | switch range { 18 | case .index: 19 | dimensionToBeRemoved.append(i) 20 | 21 | case let .slice(start: start, end: end, stride: stride): 22 | 23 | if start == 0, end == nil || end! == shape[i], stride == 1 { 24 | continue 25 | } 26 | 27 | shape[i] = Dimension(length: shape[i], memory_stride: 0) 28 | .sliced( 29 | start: start, 30 | end: end, 31 | stride: stride 32 | ) 33 | .length 34 | 35 | case let .filter(indexes): 36 | shape[i] -= indexes.count 37 | 38 | case .all: 39 | continue 40 | case .squeezeAxis: 41 | precondition( 42 | shape[i] == 1, 43 | "Cannot squeeze dimension \(i) of \(shape), expected 1 got \(shape[i])" 44 | ) 45 | 46 | dimensionToBeRemoved.append(i) 47 | 48 | case .newAxis: 49 | dimensionToBeAdded[i] = 1 50 | 51 | case .ellipsis: 52 | fatalError("Ellipsis should be expand as a series of .all expressions") 53 | } 54 | } 55 | 56 | shape = shape 57 | .enumerated() 58 | .filter { i, d in !dimensionToBeRemoved.contains(i) } 59 | .map { i, d in d } 60 | 61 | for (i, dimension) in dimensionToBeAdded { 62 | shape.insert(dimension, at: i) 63 | } 64 | 65 | return NDArray(ScalarNDArray( 66 | data, 67 | shape: shape 68 | )) 69 | } 70 | 71 | public mutating func subscript_set(_ ranges: [ArrayRange], _ ndarray: NDArray) -> NDArray { 72 | if shape.isEmpty, ranges.isEmpty { 73 | data = ndarray.scalarized() 74 | return NDArray(self) 75 | } 76 | 77 | let ndarrayView = subscript_get(ranges) 78 | 79 | if ndarrayView.shape.isEmpty, ndarrayView.shape == [1] { 80 | data = ndarray.scalarized() 81 | return NDArray(self) 82 | } 83 | 84 | return NDArray(ndarrayView.baseCopy()) 85 | } 86 | 87 | public func linearIndex(at indexes: [Int]) -> Int { 88 | 0 89 | } 90 | 91 | public func dataValue(at indexes: [Int]) -> Scalar { 92 | data 93 | } 94 | 95 | public func withScalarGetter(_ body: (@escaping NDArray.ScalarGetter) -> Void) { 96 | body { _, rectIndex in 97 | self.data 98 | } 99 | } 100 | 101 | public mutating func withScalarGetterSetter(_ body: (@escaping NDArray.ScalarGetterSetter) -> Void) { 102 | var data = self.data 103 | defer { self.data = data } 104 | 105 | body { indexer, f in 106 | data = f(data) 107 | } 108 | } 109 | 110 | public mutating func withScalarSetter(_ body: (@escaping NDArray.ScalarSetter) -> Void) { 111 | var data = self.data 112 | defer { self.data = data } 113 | 114 | body { indexer, value in 115 | data = value 116 | } 117 | } 118 | 119 | // ops 120 | public func transposed(_ indexes: [Int]) -> NDArray { 121 | NDArray(ScalarNDArray( 122 | data, 123 | shape: indexes.map { i in shape[i] } 124 | )) 125 | } 126 | 127 | public func tiled(by repetitions: [Int]) -> NDArray { 128 | NDArray(ScalarNDArray( 129 | data, 130 | shape: zip(shape, repetitions).map(*) 131 | )) 132 | } 133 | 134 | public func expandDimensions(axis: Int) -> NDArray { 135 | var shape = self.shape 136 | shape.insert(1, at: axis) 137 | 138 | return NDArray(ScalarNDArray( 139 | data, 140 | shape: shape 141 | )) 142 | } 143 | 144 | public func scalarized() -> Scalar { 145 | data 146 | } 147 | 148 | public func copy() -> NDArray { 149 | if shape.isEmpty { 150 | return NDArray(ScalarNDArray(data, shape: shape)) 151 | } else { 152 | return NDArray(baseCopy()) 153 | } 154 | } 155 | } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/add.swift: -------------------------------------------------------------------------------- 1 | 2 | extension NDArray where Scalar: AdditiveArithmetic { 3 | public static func + (lhs: NDArray, rhs: NDArray) -> NDArray { 4 | elementwise(lhs, rhs, apply: +) 5 | } 6 | 7 | public static func + (lhs: NDArray, rhs: Scalar) -> NDArray { 8 | elementwise(lhs) { $0 + rhs } 9 | } 10 | 11 | public static func + (lhs: Scalar, rhs: NDArray) -> NDArray { 12 | elementwise(rhs) { lhs + $0 } 13 | } 14 | } 15 | 16 | // import func CBlas.cblas_saxpy 17 | 18 | // extension NDArray where Scalar: AdditiveArithmetic { 19 | // @inlinable 20 | // public static func + (left: NDArray, right: Scalar) -> NDArray { 21 | // var outputData = [Scalar](repeating: right, count: left.data.count) 22 | 23 | // cblas_saxpy(Int32(left.data.count), 1, left.data, 1, &outputData, 1) 24 | 25 | // return NDArray( 26 | // outputData, 27 | // shape: left.shape 28 | // ) 29 | // } 30 | 31 | // @inlinable 32 | // public static func + (left: Scalar, right: NDArray) -> NDArray { 33 | // right + left 34 | // } 35 | 36 | // @inlinable 37 | // public static func + (left: NDArray, right: NDArray) -> NDArray { 38 | // precondition(left.shape == right.shape) 39 | 40 | // var outputData = Array(right.data) 41 | 42 | // cblas_saxpy(Int32(left.data.count), 1, left.data, 1, &outputData, 1) 43 | 44 | // return NDArray( 45 | // outputData, 46 | // shape: left.shape 47 | // ) 48 | // } 49 | // } 50 | 51 | // extension NDArray where Scalar == Float { 52 | // @inlinable 53 | // public static func + (left: NDArray, right: Scalar) -> NDArray { 54 | // var outputData = [Scalar](repeating: right, count: left.data.count) 55 | 56 | // cblas_saxpy(Int32(left.data.count), 1, left.data, 1, &outputData, 1) 57 | 58 | // return NDArray( 59 | // outputData, 60 | // shape: left.shape 61 | // ) 62 | // } 63 | 64 | // @inlinable 65 | // public static func + (left: NDArray, right: NDArray) -> NDArray { 66 | // precondition(left.shape == right.shape) 67 | 68 | // var outputData = Array(right.data) 69 | 70 | // cblas_saxpy(Int32(left.data.count), 1, left.data, 1, &outputData, 1) 71 | 72 | // return NDArray( 73 | // outputData, 74 | // shape: left.shape 75 | // ) 76 | // } 77 | // } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/broadcast.swift: -------------------------------------------------------------------------------- 1 | 2 | @usableFromInline 3 | internal func broadcast(_ left: NDArray, and right: NDArray) -> (left: NDArray, right: NDArray) { 4 | if left.shape == right.shape { 5 | return ( 6 | left: left, 7 | right: right 8 | ) 9 | } 10 | 11 | var left = left 12 | var right = right 13 | 14 | precondition( 15 | left.shape == [] || right.shape == [] || zip(left.shape, right.shape).map { leftLength, rightLength in 16 | leftLength == rightLength || leftLength == 1 || rightLength == 1 17 | }.reduce(true) { $0 && $1 }, 18 | "Cannot broadcast shapes \(left.shape) and \(right.shape)" 19 | ) 20 | 21 | if left.shape.count == 0 { 22 | for _ in 0 ..< right.shape.count { 23 | left = left.expandDimensions(axis: 0) 24 | } 25 | } else if right.shape.count == 0 { 26 | for _ in 0 ..< left.shape.count { 27 | right = right.expandDimensions(axis: 0) 28 | } 29 | } 30 | 31 | var leftRepetitions = Array(repeating: 1, count: left.shape.count) 32 | var rightRepetitions = Array(repeating: 1, count: left.shape.count) 33 | 34 | for i in 0 ..< left.shape.count { 35 | if left.shape[i] == right.shape[i] { 36 | continue 37 | } else if left.shape[i] == 1 { 38 | leftRepetitions[i] = right.shape[i] 39 | } else { 40 | rightRepetitions[i] = left.shape[i] 41 | } 42 | } 43 | 44 | left = left.tiled(by: leftRepetitions) 45 | right = right.tiled(by: rightRepetitions) 46 | 47 | return ( 48 | left: left, 49 | right: right 50 | ) 51 | } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/description.swift: -------------------------------------------------------------------------------- 1 | 2 | extension NDArray: CustomStringConvertible { 3 | public var description: String { 4 | let nElements = shape.product() 5 | var s = "\(Self.self)\(shape)(" + String(repeating: "[", count: Swift.max(shape.count - 1, 0)) 6 | 7 | withScalarGetter { valueAt in 8 | 9 | if self.shape.count == 0 { 10 | s += "\(self.dataValue(at: [Int]()))" + ")" 11 | } else if self.shape.count == 1 { 12 | var arrayString = "" 13 | for index in indexSequence(range: 0 ..< nElements, shape: self.shape) { 14 | let (i, _) = index 15 | arrayString += "\(valueAt(index))" + (i + 1 != nElements ? ", " : "") 16 | } 17 | s += "[\(arrayString)])" 18 | } else { 19 | let reversedShape = Array(self.shape.reversed()) 20 | let lastDim = reversedShape[0] 21 | let secondLastDim = reversedShape[1] 22 | 23 | s += "\n" 24 | 25 | var arrayString = "" 26 | 27 | for index in indexSequence(range: 0 ..< nElements + 1, shape: self.shape) { 28 | let (i, _) = index 29 | 30 | if i % lastDim == 0, i > 0 { 31 | s += " [\(arrayString)],\n" 32 | arrayString = "" 33 | } 34 | 35 | if i % (lastDim * secondLastDim) == 0, i > 0, i < nElements { 36 | s += "\n" 37 | } 38 | 39 | if i < nElements { 40 | arrayString += "\(valueAt(index))" + ((i + 1) % lastDim != 0 ? ", " : "") 41 | } 42 | } 43 | s += String(repeating: "]", count: Swift.max(self.shape.count - 1, 1)) + ")" 44 | } 45 | } 46 | 47 | return s 48 | } 49 | } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/div.swift: -------------------------------------------------------------------------------- 1 | 2 | extension NDArray where Scalar: Divisible { 3 | public static func / (lhs: NDArray, rhs: NDArray) -> NDArray { 4 | elementwise(lhs, rhs, apply: /) 5 | } 6 | 7 | public static func / (lhs: Scalar, rhs: NDArray) -> NDArray { 8 | elementwise(rhs) { lhs / $0 } 9 | } 10 | 11 | public static func / (lhs: NDArray, rhs: Scalar) -> NDArray { 12 | elementwise(lhs) { $0 / rhs } 13 | } 14 | } 15 | 16 | // import func CBlas.cblas_saxpy 17 | // import func CBlas.cblas_sscal 18 | 19 | // extension NDArray where Scalar == Float { 20 | // public static func / (_ left: NDArray, _ right: Scalar) -> NDArray { 21 | // let right = 1 / right 22 | // var outputData = Array(left.data) 23 | 24 | // cblas_sscal(Int32(left.data.count), right, &outputData, 1) 25 | 26 | // return NDArray( 27 | // outputData, 28 | // shape: left.shape 29 | // ) 30 | // } 31 | 32 | // public static func / (_ left: Scalar, _ right: NDArray) -> NDArray { 33 | // let left = NDArray( 34 | // Array(repeating: left, count: right.shape.product()), 35 | // shape: right.shape 36 | // ) 37 | 38 | // return left / right 39 | // } 40 | 41 | // public static func / (_ left: NDArray, _ right: NDArray) -> NDArray { 42 | // precondition(left.shape == right.shape) 43 | 44 | // return elementwise( 45 | // left, 46 | // right, 47 | // apply: / 48 | // ) 49 | // } 50 | // } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/dot.swift: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cgarciae/NDArray/49c5f1ec9958a73afde90842c57995d8de9f57b7/Sources/NDArray/Ops/dot.swift -------------------------------------------------------------------------------- /Sources/NDArray/Ops/elementWise.swift: -------------------------------------------------------------------------------- 1 | 2 | import Foundation 3 | 4 | public let DISPATCH = DispatchQueue(label: "NDArray", attributes: .concurrent) 5 | public let CPU_COUNT = ProcessInfo.processInfo.activeProcessorCount 6 | 7 | ////////////////////////////////////////////////////////////////////////////////////////// 8 | // 1 9 | ////////////////////////////////////////////////////////////////////////////////////////// 10 | 11 | // @inlinable public func getIndexer(_ ndarray: NDArray) -> ((Int, UnsafeMutableBufferPointer)) -> Int { 12 | // ndarray.arrayShape.isOriginalShape ? 13 | // { $0.0 } : { ndarray.arrayShape.linearIndex(of: $0.1) } 14 | // } 15 | 16 | @inlinable 17 | public func elementwise( 18 | _ ndArrayA: NDArray, 19 | into ndArrayZ: inout NDArray, 20 | apply f: (A) -> Z 21 | ) { 22 | let nElements = ndArrayA.shape.product() 23 | 24 | ndArrayZ.withScalarSetter { zSetter in 25 | ndArrayA.withScalarGetter { aGetter in 26 | 27 | for index in indexSequence(range: 0 ..< nElements, shape: ndArrayA.shape) { 28 | let value = f(aGetter(index)) 29 | zSetter(index, value) 30 | } 31 | } 32 | } 33 | } 34 | 35 | @inlinable 36 | public func elementwise( 37 | _ ndArrayA: NDArray, 38 | apply f: (A) -> Z 39 | ) -> NDArray { 40 | let nElements = ndArrayA.shape.product() 41 | 42 | var ndArrayZ = NDArray( 43 | [Z](unsafeUninitializedCapacity: nElements) { x, count in 44 | count = nElements 45 | }, 46 | shape: ndArrayA.shape 47 | ) 48 | 49 | elementwise(ndArrayA, into: &ndArrayZ, apply: f) 50 | 51 | return ndArrayZ 52 | } 53 | 54 | @inlinable 55 | public func elementwiseInParallel( 56 | _ ndArrayA: NDArray, 57 | into ndArrayZ: inout NDArray, 58 | workers: Int = CPU_COUNT, 59 | apply f: @escaping (A) -> Z 60 | ) { 61 | let nElements = ndArrayA.shape.product() 62 | 63 | ndArrayZ.withScalarSetter { zSetter in 64 | ndArrayA.withScalarGetter { aGetter in 65 | 66 | let rangeMap = { indexSequence(range: $0, shape: ndArrayA.shape) } 67 | 68 | parFor(0 ..< nElements, rangeMap: rangeMap) { index in 69 | let value = f(aGetter(index)) 70 | zSetter(index, value) 71 | } 72 | } 73 | } 74 | } 75 | 76 | @inlinable 77 | public func elementwiseInParallel( 78 | _ ndArrayA: NDArray, 79 | workers: Int = CPU_COUNT, 80 | apply f: @escaping (A) -> Z 81 | ) -> NDArray { 82 | let nElements = ndArrayA.shape.product() 83 | 84 | var ndArrayZ = NDArray( 85 | [Z](unsafeUninitializedCapacity: nElements) { x, count in 86 | count = nElements 87 | }, 88 | shape: ndArrayA.shape 89 | ) 90 | 91 | elementwiseInParallel(ndArrayA, into: &ndArrayZ, workers: workers, apply: f) 92 | 93 | return ndArrayZ 94 | } 95 | 96 | ////////////////////////////////////////////////////////////////////////////////////////// 97 | // 2 98 | ////////////////////////////////////////////////////////////////////////////////////////// 99 | 100 | @inlinable 101 | public func elementwise( 102 | _ ndArrayA: NDArray, 103 | _ ndArrayB: NDArray, 104 | into ndArrayZ: inout NDArray, 105 | apply f: (A, B) -> Z 106 | ) { 107 | var ndArrayA = ndArrayA 108 | var ndArrayB = ndArrayB 109 | 110 | if ndArrayA.shape != ndArrayB.shape { 111 | (ndArrayA, ndArrayB) = broadcast(ndArrayA, and: ndArrayB) 112 | } 113 | 114 | let nElements = ndArrayA.shape.product() 115 | 116 | ndArrayZ.withScalarSetter { zSetter in 117 | ndArrayA.withScalarGetter { aGetter in 118 | ndArrayB.withScalarGetter { bGetter in 119 | 120 | for index in indexSequence(range: 0 ..< nElements, shape: ndArrayA.shape) { 121 | let value = f(aGetter(index), bGetter(index)) 122 | zSetter(index, value) 123 | } 124 | } 125 | } 126 | } 127 | } 128 | 129 | @inlinable 130 | public func elementwiseAssignApply( 131 | _ ndArrayA: inout NDArray, 132 | _ ndArrayB: NDArray, 133 | apply f: (A, B) -> A 134 | ) { 135 | var ndArrayB = ndArrayB 136 | 137 | if ndArrayA.shape != ndArrayB.shape { 138 | (ndArrayA, ndArrayB) = broadcast(ndArrayA, and: ndArrayB) 139 | } 140 | 141 | let nElements = ndArrayA.shape.product() 142 | 143 | let ndArrayAShape = ndArrayA.shape 144 | 145 | ndArrayA.withScalarGetterSetter { aGetterSetter in 146 | ndArrayB.withScalarGetter { bGetter in 147 | 148 | for index in indexSequence(range: 0 ..< nElements, shape: ndArrayAShape) { 149 | aGetterSetter(index) { aValue in 150 | let value = f(aValue, bGetter(index)) 151 | return value 152 | } 153 | } 154 | } 155 | } 156 | } 157 | 158 | @inlinable 159 | public func elementwise( 160 | _ ndArrayA: NDArray, 161 | _ ndArrayB: NDArray, 162 | apply f: (A, B) -> Z 163 | ) -> NDArray { 164 | var ndArrayA = ndArrayA 165 | var ndArrayB = ndArrayB 166 | 167 | if ndArrayA.shape != ndArrayB.shape { 168 | (ndArrayA, ndArrayB) = broadcast(ndArrayA, and: ndArrayB) 169 | } 170 | 171 | let nElements = ndArrayA.shape.product() 172 | 173 | var ndArrayZ = NDArray( 174 | [Z](unsafeUninitializedCapacity: nElements) { x, count in 175 | count = nElements 176 | }, 177 | shape: ndArrayA.shape 178 | ) 179 | 180 | elementwise(ndArrayA, ndArrayB, into: &ndArrayZ, apply: f) 181 | 182 | return ndArrayZ 183 | } 184 | 185 | @inlinable 186 | public func elementwiseInParallel( 187 | _ ndArrayA: NDArray, 188 | _ ndArrayB: NDArray, 189 | into ndArrayZ: inout NDArray, 190 | workers: Int = CPU_COUNT, 191 | apply f: @escaping (A, B) -> Z 192 | ) { 193 | precondition(ndArrayA.shape == ndArrayB.shape) 194 | let nElements = ndArrayA.shape.product() 195 | 196 | if !ndArrayZ.anyNDArray.isSetable() { 197 | ndArrayZ = NDArray(ndArrayZ.baseCopy()) 198 | } 199 | 200 | ndArrayZ.withScalarSetter { zSetter in 201 | ndArrayA.withScalarGetter { aGetter in 202 | ndArrayB.withScalarGetter { bGetter in 203 | 204 | let rangeMap = { indexSequence(range: $0, shape: ndArrayA.shape) } 205 | 206 | parFor(0 ..< nElements, rangeMap: rangeMap) { index in 207 | 208 | let value = f(aGetter(index), bGetter(index)) 209 | zSetter(index, value) 210 | } 211 | } 212 | } 213 | } 214 | } 215 | 216 | @inlinable 217 | public func elementwiseInParallel( 218 | _ ndArrayA: NDArray, 219 | _ ndArrayB: NDArray, 220 | workers: Int = CPU_COUNT, 221 | apply f: @escaping (A, B) -> Z 222 | ) -> NDArray { 223 | var ndArrayA = ndArrayA 224 | var ndArrayB = ndArrayB 225 | 226 | if ndArrayA.shape != ndArrayB.shape { 227 | (ndArrayA, ndArrayB) = broadcast(ndArrayA, and: ndArrayB) 228 | } 229 | 230 | let nElements = ndArrayA.shape.product() 231 | 232 | var ndArrayZ = NDArray( 233 | [Z](unsafeUninitializedCapacity: nElements) { x, count in 234 | count = nElements 235 | }, 236 | shape: ndArrayA.shape 237 | ) 238 | 239 | elementwiseInParallel(ndArrayA, ndArrayB, into: &ndArrayZ, workers: workers, apply: f) 240 | 241 | return ndArrayZ 242 | } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/max.swift: -------------------------------------------------------------------------------- 1 | 2 | extension NDArray where Scalar: AdditiveArithmetic & Comparable { 3 | public func max(axis: Int) -> NDArray { 4 | max(axis: [axis]) 5 | } 6 | 7 | public func max(axis: [Int]) -> NDArray { 8 | let reducingShape = shape 9 | .enumerated() 10 | .filter { axis.contains($0.offset) } 11 | .map { $0.element } 12 | 13 | let nonReducedAxis = Array( 14 | Set(0 ..< shape.count).subtracting(Set(axis)) 15 | ) 16 | 17 | var ranges = reducingShape.map { _ in ArrayRange.index(0) } 18 | 19 | for i in nonReducedAxis.sorted() { 20 | ranges.insert(.all, at: i) 21 | } 22 | 23 | // replace with NDArray(zeros: ...) in the future 24 | var initial = self[r: ranges].copy() 25 | 26 | return reduce(axis: axis, initial: &initial) { 27 | Swift.max($0, $1) 28 | } 29 | } 30 | } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/mean.swift: -------------------------------------------------------------------------------- 1 | 2 | 3 | extension NDArray where Scalar: FloatingPoint & Numeric { 4 | public func mean(axis: Int) -> NDArray { 5 | mean(axis: [axis]) 6 | } 7 | 8 | public func mean(axis: [Int]) -> NDArray { 9 | let outputShape = shape 10 | .enumerated() 11 | .filter { !axis.contains($0.offset) } 12 | .map { $0.element } 13 | 14 | let reductionElements = Scalar(shape 15 | .enumerated() 16 | .filter { axis.contains($0.offset) } 17 | .map { $0.element } 18 | .product() 19 | ) 20 | 21 | // replace with NDArray(zeros: ...) in the future 22 | var initial = NDArray( 23 | [Scalar]( 24 | repeating: Scalar.zero, 25 | count: outputShape.product() 26 | ), 27 | shape: outputShape 28 | ) 29 | 30 | return reduce(axis: axis, initial: &initial) { acc, x in 31 | acc + x / reductionElements 32 | } 33 | } 34 | } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/min.swift: -------------------------------------------------------------------------------- 1 | 2 | extension NDArray where Scalar: AdditiveArithmetic & Comparable { 3 | public func min(axis: Int) -> NDArray { 4 | min(axis: [axis]) 5 | } 6 | 7 | public func min(axis: [Int]) -> NDArray { 8 | let reducingShape = shape 9 | .enumerated() 10 | .filter { axis.contains($0.offset) } 11 | .map { $0.element } 12 | 13 | let nonReducedAxis = Array( 14 | Set(0 ..< shape.count).subtracting(Set(axis)) 15 | ) 16 | 17 | var ranges = reducingShape.map { _ in ArrayRange.index(0) } 18 | 19 | for i in nonReducedAxis.sorted() { 20 | ranges.insert(.all, at: i) 21 | } 22 | 23 | // replace with NDArray(zeros: ...) in the future 24 | var initial = self[r: ranges].copy() 25 | 26 | return reduce(axis: axis, initial: &initial) { 27 | Swift.min($0, $1) 28 | } 29 | } 30 | } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/mul.swift: -------------------------------------------------------------------------------- 1 | 2 | extension NDArray where Scalar: Numeric { 3 | public static func * (lhs: NDArray, rhs: NDArray) -> NDArray { 4 | elementwise(lhs, rhs, apply: *) 5 | } 6 | 7 | public static func * (lhs: Scalar, rhs: NDArray) -> NDArray { 8 | elementwise(rhs) { lhs * $0 } 9 | } 10 | 11 | public static func * (lhs: NDArray, rhs: Scalar) -> NDArray { 12 | elementwise(lhs) { $0 * rhs } 13 | } 14 | } 15 | 16 | // import func CBlas.cblas_saxpy 17 | // import func CBlas.cblas_sscal 18 | 19 | // extension NDArray where Scalar == Float { 20 | // public static func * (_ left: NDArray, _ right: Float) -> NDArray { 21 | // var outputData = Array(left.data) 22 | 23 | // cblas_sscal(Int32(left.data.count), right, &outputData, 1) 24 | 25 | // return NDArray( 26 | // outputData, 27 | // shape: left.shape 28 | // ) 29 | // } 30 | 31 | // public static func * (_ left: Float, _ right: NDArray) -> NDArray { 32 | // right * left 33 | // } 34 | 35 | // public static func * (_ left: NDArray, _ right: NDArray) -> NDArray { 36 | // precondition(left.shape == right.shape) 37 | 38 | // return NDArray( 39 | // elementwise( 40 | // left.data, 41 | // right.data, 42 | // apply: * 43 | // ), 44 | // shape: left.shape 45 | // ) 46 | // } 47 | // } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/reduce.swift: -------------------------------------------------------------------------------- 1 | 2 | 3 | extension NDArray { 4 | public func reduce(axis: [Int], initial: NDArray, f: (B, Scalar) -> B) -> NDArray { 5 | var initial = initial 6 | return reduce(axis: axis, initial: &initial, f: f) 7 | } 8 | 9 | public func reduce(axis: [Int], initial: inout NDArray, f: (B, Scalar) -> B) -> NDArray { 10 | if !initial.anyNDArray.isSetable() { 11 | initial = NDArray(initial.baseCopy()) 12 | } 13 | 14 | let reducingShape = shape 15 | .enumerated() 16 | .filter { axis.contains($0.offset) } 17 | .map { $0.element } 18 | 19 | let nonReducedAxis = Array( 20 | Set(0 ..< shape.count).subtracting(Set(axis)) 21 | ) 22 | 23 | for index in indexSequence(range: 0 ..< reducingShape.product(), shape:reducingShape) { 24 | var ranges = index.rectangularIndex.map { ArrayRange.index($0) } 25 | 26 | for i in nonReducedAxis.sorted() { 27 | ranges.insert(.all, at: i) 28 | } 29 | 30 | let view = self[r: ranges] 31 | 32 | elementwiseAssignApply(&initial, view, apply: f) 33 | } 34 | 35 | return initial 36 | } 37 | } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/sub.swift: -------------------------------------------------------------------------------- 1 | extension NDArray where Scalar: AdditiveArithmetic { 2 | public static func - (lhs: NDArray, rhs: NDArray) -> NDArray { 3 | elementwise(lhs, rhs, apply: -) 4 | } 5 | 6 | public static func - (lhs: Scalar, rhs: NDArray) -> NDArray { 7 | elementwise(rhs) { lhs - $0 } 8 | } 9 | 10 | public static func - (lhs: NDArray, rhs: Scalar) -> NDArray { 11 | elementwise(lhs) { $0 - rhs } 12 | } 13 | } 14 | 15 | // import func CBlas.cblas_saxpy 16 | 17 | // extension NDArray where Scalar == Float { 18 | // @inlinable 19 | // public static func - (left: NDArray, right: Scalar) -> NDArray { 20 | // var outputData = [Scalar](repeating: -right, count: left.data.count) 21 | 22 | // cblas_saxpy(Int32(left.data.count), 1, left.data, 1, &outputData, 1) 23 | 24 | // return NDArray( 25 | // outputData, 26 | // shape: left.shape 27 | // ) 28 | // } 29 | 30 | // @inlinable 31 | // public static func - (left: Scalar, right: NDArray) -> NDArray { 32 | // right - left 33 | // } 34 | 35 | // @inlinable 36 | // public static func - (left: NDArray, right: NDArray) -> NDArray { 37 | // precondition(left.shape == right.shape) 38 | 39 | // var outputData = Array(left.data) 40 | 41 | // cblas_saxpy(Int32(left.data.count), -1, right.data, 1, &outputData, 1) 42 | 43 | // return NDArray( 44 | // outputData, 45 | // shape: left.shape 46 | // ) 47 | // } 48 | // } 49 | 50 | // internal extension NDArray where Scalar == Float { 51 | // @inlinable 52 | // static func _vjpSubtract(lhs: NDArray, rhs: NDArray) -> 53 | // ( 54 | // NDArray, 55 | // (NDArray) -> (NDArray, NDArray) 56 | // ) { 57 | // let value = lhs - rhs 58 | // return ( 59 | // value, 60 | // { dvalue in (dvalue * lhs, -1 * dvalue * rhs) } 61 | // ) 62 | // } 63 | // } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/subscript.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | 3 | public let all = ArrayRange.all 4 | public let newAxis = ArrayRange.newAxis 5 | public let ellipsis = ArrayRange.ellipsis 6 | public let squeezeAxis = ArrayRange.squeezeAxis 7 | 8 | // experimental alias 9 | public let rest = ArrayRange.ellipsis 10 | public let new = ArrayRange.newAxis 11 | public let squeeze = ArrayRange.squeezeAxis 12 | 13 | extension NDArray { 14 | @inlinable 15 | public subscript(_ ranges: ArrayExpression...) -> NDArray { 16 | get { 17 | self[r: ranges.map { $0.arrayRange }] 18 | } 19 | mutating set(ndarray) { 20 | self[r: ranges.map { $0.arrayRange }] = ndarray 21 | } 22 | } 23 | 24 | @inlinable 25 | public subscript(r ranges: [ArrayRange]) -> NDArray { 26 | get { 27 | subscript_get(ranges) 28 | } 29 | 30 | mutating set(ndarray) { 31 | self = subscript_set(ranges, ndarray) 32 | } 33 | } 34 | } 35 | 36 | public protocol ArrayExpression { 37 | @inlinable 38 | var arrayRange: ArrayRange { get } 39 | } 40 | 41 | extension Array: ArrayExpression { 42 | public var arrayRange: ArrayRange { 43 | if self is [Int] { 44 | return .filter(self as! [Int]) 45 | } else if self is [Bool] { 46 | let array = self as! [Bool] 47 | return .filter(array.enumerated().filter { $0.1 }.map { $0.0 }) 48 | } else { 49 | fatalError("Type \(Element.self) not supported") 50 | } 51 | } 52 | } 53 | 54 | public enum ArrayRange: ArrayExpression { 55 | case ellipsis 56 | case newAxis 57 | case squeezeAxis 58 | case all 59 | case index(Int) 60 | case filter([Int]) 61 | case slice(start: Int? = nil, end: Int? = nil, stride: Int = 1) 62 | 63 | public var arrayRange: ArrayRange { self } 64 | } 65 | 66 | public func isEllipsis(_ range: ArrayRange) -> Bool { 67 | switch range { 68 | case .ellipsis: 69 | return true 70 | default: 71 | return false 72 | } 73 | } 74 | 75 | extension Int: ArrayExpression { 76 | public var arrayRange: ArrayRange { .index(self) } 77 | } 78 | 79 | extension Range: ArrayExpression where Bound == Int { 80 | public var arrayRange: ArrayRange { .slice(start: lowerBound, end: upperBound) } 81 | public func stride(_ stride: Int) -> ArrayRange { 82 | .slice(start: lowerBound, end: upperBound, stride: stride) 83 | } 84 | } 85 | 86 | extension ClosedRange: ArrayExpression where Bound == Int { 87 | public var arrayRange: ArrayRange { .slice(start: lowerBound, end: upperBound + 1) } 88 | public func stride(_ stride: Int) -> ArrayRange { 89 | .slice(start: lowerBound, end: upperBound + 1, stride: stride) 90 | } 91 | } 92 | 93 | extension PartialRangeFrom: ArrayExpression where Bound == Int { 94 | public var arrayRange: ArrayRange { .slice(start: lowerBound) } 95 | public func stride(_ stride: Int) -> ArrayRange { 96 | .slice(start: lowerBound, stride: stride) 97 | } 98 | } 99 | 100 | extension PartialRangeUpTo: ArrayExpression where Bound == Int { 101 | public var arrayRange: ArrayRange { .slice(end: upperBound) } 102 | public func stride(_ stride: Int) -> ArrayRange { 103 | .slice(end: upperBound, stride: stride) 104 | } 105 | } 106 | 107 | extension PartialRangeThrough: ArrayExpression where Bound == Int { 108 | public var arrayRange: ArrayRange { .slice(end: upperBound + 1) } 109 | public func stride(_ stride: Int) -> ArrayRange { 110 | .slice(end: upperBound + 1, stride: stride) 111 | } 112 | } 113 | 114 | public struct Slice: ArrayExpression { 115 | let start: Int? 116 | let end: Int? 117 | 118 | internal init(start: Int? = nil, end: Int? = nil) { 119 | self.start = start 120 | self.end = end 121 | } 122 | 123 | public var arrayRange: ArrayRange { 124 | .slice(start: start, end: end) 125 | } 126 | } 127 | 128 | public struct StridedSlice : ArrayExpression { 129 | let start: Int? 130 | let end: Int? 131 | let stride: Int 132 | 133 | internal init(start: Int? = nil, end: Int? = nil, stride: Int = 1) { 134 | self.start = start 135 | self.end = end 136 | self.stride = stride 137 | } 138 | 139 | public var arrayRange: ArrayRange { 140 | .slice(start: start, end: end, stride: stride) 141 | } 142 | } 143 | 144 | ///////////////////////////////////////////////////////////////////////////////////////////// 145 | // .. 146 | ///////////////////////////////////////////////////////////////////////////////////////////// 147 | 148 | postfix operator .. 149 | prefix operator .. 150 | prefix operator ..- 151 | infix operator ..: AdditionPrecedence 152 | infix operator ..-: AdditionPrecedence 153 | infix operator .... 154 | prefix operator .... 155 | prefix operator ....- 156 | 157 | public extension Int { 158 | static postfix func .. (lhs: Int) -> Slice { 159 | Slice(start: lhs) 160 | } 161 | 162 | static prefix func .. (rhs: Int) -> Slice { 163 | Slice(end: rhs) 164 | } 165 | 166 | static prefix func ..- (rhs: Int) -> Slice { 167 | Slice(end: -rhs) 168 | } 169 | 170 | static func .. (lhs: Int, rhs: Int) -> Slice { 171 | Slice(start: lhs, end: rhs) 172 | } 173 | 174 | static func .. (lhs: Slice, rhs: Int) -> StridedSlice { 175 | StridedSlice(start: lhs.start, end: lhs.end, stride: rhs) 176 | } 177 | 178 | static func ..- (lhs: Int, rhs: Int) -> Slice { 179 | Slice(start: lhs, end: -rhs) 180 | } 181 | 182 | static func ..- (lhs: Slice, rhs: Int) -> StridedSlice { 183 | StridedSlice(start: lhs.start, end: lhs.end, stride: -rhs) 184 | } 185 | 186 | static prefix func .... (rhs: Int) -> StridedSlice { 187 | StridedSlice(stride: rhs) 188 | } 189 | 190 | static prefix func ....- (rhs: Int) -> StridedSlice { 191 | StridedSlice(stride: -rhs) 192 | } 193 | 194 | static func .... (lhs: Int, rhs: Int) -> StridedSlice { 195 | StridedSlice(start: lhs, stride: rhs) 196 | } 197 | } 198 | 199 | ///////////////////////////////////////////////////////////////////////////////////////////// 200 | // |> 201 | ///////////////////////////////////////////////////////////////////////////////////////////// 202 | 203 | // postfix operator |> 204 | // prefix operator |> 205 | // prefix operator |>- 206 | // infix operator |>: AdditionPrecedence 207 | // infix operator ||> 208 | // prefix operator ||> 209 | // prefix operator ||>- 210 | 211 | // public extension Int { 212 | // static postfix func |> (lhs: Int) -> Slice { 213 | // Slice(start: lhs) 214 | // } 215 | 216 | // static prefix func |> (rhs: Int) -> Slice { 217 | // Slice(end: rhs) 218 | // } 219 | 220 | // static prefix func |>- (rhs: Int) -> Slice { 221 | // Slice(end: -rhs) 222 | // } 223 | 224 | // static func |> (lhs: Int, rhs: Int) -> Slice { 225 | // Slice(start: lhs, end: rhs) 226 | // } 227 | 228 | // static func |> (lhs: Slice, rhs: Int) -> Slice { 229 | // Slice(start: lhs.start, end: lhs.end, stride: rhs) 230 | // } 231 | 232 | // static func |> (lhs: Int, rhs: Slice) -> Slice { 233 | // Slice(start: lhs, end: rhs.start, stride: rhs.end!) 234 | // } 235 | 236 | // static prefix func ||> (rhs: Int) -> Slice { 237 | // Slice(stride: rhs) 238 | // } 239 | 240 | // static prefix func ||>- (rhs: Int) -> Slice { 241 | // Slice(stride: -rhs) 242 | // } 243 | 244 | // static func ||> (lhs: Int, rhs: Int) -> Slice { 245 | // Slice(start: lhs, stride: rhs) 246 | // } 247 | // } -------------------------------------------------------------------------------- /Sources/NDArray/Ops/sum.swift: -------------------------------------------------------------------------------- 1 | 2 | extension NDArray where Scalar: AdditiveArithmetic { 3 | public func sum(axis: Int) -> NDArray { 4 | sum(axis: [axis]) 5 | } 6 | 7 | public func sum(axis: [Int]) -> NDArray { 8 | let outputShape = shape 9 | .enumerated() 10 | .filter { !axis.contains($0.offset) } 11 | .map { $0.element } 12 | 13 | // replace with NDArray(zeros: ...) in the future 14 | var initial = NDArray( 15 | [Scalar]( 16 | repeating: Scalar.zero, 17 | count: outputShape.product() 18 | ), 19 | shape: outputShape 20 | ) 21 | 22 | return reduce(axis: axis, initial: &initial, f: +) 23 | } 24 | } -------------------------------------------------------------------------------- /Sources/NDArray/Protocols/Divisible.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | 3 | public protocol Divisible: Numeric { 4 | static func / (lhs: Self, rhs: Self) -> Self 5 | static func /= (lhs: inout Self, rhs: Self) 6 | } 7 | 8 | extension Divisible { 9 | public static func /= (lhs: inout Self, rhs: Self) { 10 | lhs = lhs / rhs 11 | } 12 | } 13 | 14 | extension Int: Divisible {} 15 | extension Int8: Divisible {} 16 | extension Int16: Divisible {} 17 | extension Int32: Divisible {} 18 | extension Int64: Divisible {} 19 | 20 | extension UInt: Divisible {} 21 | extension UInt8: Divisible {} 22 | extension UInt16: Divisible {} 23 | extension UInt32: Divisible {} 24 | extension UInt64: Divisible {} 25 | 26 | extension Double: Divisible {} 27 | extension Float: Divisible {} 28 | extension Float80: Divisible {} 29 | extension CGFloat: Divisible {} -------------------------------------------------------------------------------- /Sources/NDArray/Protocols/NDArrayFloatingPoint.swift: -------------------------------------------------------------------------------- 1 | public protocol NDArrayFloatingPoint: Numeric & Divisible 2 | // : 3 | // BinaryFloatingPoint & Differentiable & ElementaryFunctions 4 | // where Self.RawSignificand: FixedWidthInteger, 5 | // Self == Self.TangentVector, 6 | // Self == Self.AllDifferentiableVariables 7 | {} 8 | 9 | extension Float: NDArrayFloatingPoint {} 10 | extension Double: NDArrayFloatingPoint {} -------------------------------------------------------------------------------- /Sources/NDArray/Protocols/NDArrayProtocol.swift: -------------------------------------------------------------------------------- 1 | 2 | 3 | public protocol NDArrayProtocol { 4 | associatedtype Scalar 5 | 6 | typealias ScalarGetter = ((Int, UnsafeMutableBufferPointer)) -> Scalar 7 | typealias ScalarSetter = ((Int, UnsafeMutableBufferPointer), Scalar) -> Void 8 | typealias ScalarGetterSetter = ((Int, UnsafeMutableBufferPointer), (Scalar) -> Scalar) -> Void 9 | 10 | var shape: [Int] { get } 11 | 12 | func subscript_get(_: [ArrayRange]) -> NDArray 13 | mutating func subscript_set(_: [ArrayRange], _: NDArray) -> NDArray 14 | func linearIndex(at indexes: [Int]) -> Int 15 | func dataValue(at indexes: [Int]) -> Scalar 16 | func withScalarGetter(_: (@escaping ScalarGetter) -> Void) 17 | mutating func withScalarSetter(_: (@escaping ScalarSetter) -> Void) 18 | mutating func withScalarGetterSetter(_: (@escaping ScalarGetterSetter) -> Void) 19 | 20 | // ops 21 | func transposed(_: [Int]) -> NDArray 22 | func tiled(by: [Int]) -> NDArray 23 | func expandDimensions(axis: Int) -> NDArray 24 | func scalarized() -> Scalar 25 | 26 | // 27 | // mutating func copyInternals() -> Void 28 | func copy() -> NDArray 29 | } 30 | 31 | public protocol MultiArray { 32 | associatedtype Element 33 | 34 | var array: [Element] { get } 35 | } 36 | 37 | extension Array : MultiArray { 38 | public var array: [Element] { self } 39 | } 40 | 41 | public protocol SetableNDArray {} -------------------------------------------------------------------------------- /Sources/NDArray/Ref.swift: -------------------------------------------------------------------------------- 1 | 2 | 3 | public final class Ref: CustomStringConvertible { 4 | @usableFromInline var value: A 5 | 6 | @usableFromInline init(_ value: A) { 7 | self.value = value 8 | } 9 | 10 | public var description: String { "\(value)" } 11 | } -------------------------------------------------------------------------------- /Sources/NDArray/utils.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | 3 | public func indexSequence(range: Range, shape: [Int]) -> AnySequence<(linearIndex: Int, rectangularIndex: UnsafeMutableBufferPointer)> { 4 | AnySequence { () -> AnyIterator<(linearIndex: Int, rectangularIndex: UnsafeMutableBufferPointer)> in 5 | let arrayShape = shape 6 | 7 | var iterator = range.makeIterator() 8 | let dimensionStrides = getDimensionStrides(of: arrayShape) 9 | 10 | let rectangularIndex = UnsafeMutableBufferPointer.allocate(capacity: arrayShape.count) 11 | rectangularIndex.initialize(repeating: 0) 12 | let shape = UnsafeMutableBufferPointer.allocate(capacity: arrayShape.count) 13 | _ = shape.initialize(from: arrayShape) 14 | 15 | var first = true 16 | return AnyIterator { () -> (linearIndex: Int, rectangularIndex: UnsafeMutableBufferPointer)? in 17 | guard let current = iterator.next() else { 18 | shape.baseAddress!.deinitialize(count: shape.count) 19 | shape.deallocate() 20 | 21 | rectangularIndex.baseAddress!.deinitialize(count: rectangularIndex.count) 22 | rectangularIndex.deallocate() 23 | 24 | return nil 25 | } 26 | 27 | if first { 28 | first = false 29 | var remainder = current 30 | 31 | for i in 0 ..< shape.count { 32 | if shape[i] > 1 { 33 | let index: Int 34 | (index, remainder) = remainder.quotientAndRemainder(dividingBy: dimensionStrides[i]) 35 | 36 | rectangularIndex[i] = index 37 | } 38 | } 39 | 40 | return ( 41 | linearIndex: current, 42 | rectangularIndex: rectangularIndex 43 | ) 44 | } 45 | 46 | var pos = shape.count - 1 47 | 48 | while pos >= 0 { 49 | let nextValue = rectangularIndex[pos] + 1 50 | if nextValue % shape[pos] == 0 { 51 | rectangularIndex[pos] = 0 52 | pos -= 1 53 | } else { 54 | rectangularIndex[pos] = nextValue 55 | break 56 | } 57 | } 58 | 59 | return ( 60 | linearIndex: current, 61 | rectangularIndex: rectangularIndex 62 | ) 63 | } 64 | } 65 | } 66 | 67 | extension Array { 68 | func chunked(into size: Int) -> [[Element]] { 69 | return stride(from: 0, to: count, by: size).map { 70 | Array(self[$0 ..< $0 + size]) 71 | } 72 | } 73 | } 74 | 75 | public func flattenArrays(_ array: [Any]) -> (array: [A], shape: [Int]) { 76 | if array.count == 0 { 77 | return ( 78 | array: [], 79 | shape: [] 80 | ) 81 | } else { 82 | var shape: [Int] = [array.count] 83 | return flattenArrays(array: array, shape: &shape) 84 | } 85 | } 86 | 87 | func flattenArrays(array: [Any], shape: inout [Int]) -> (array: [A], shape: [Int]) { 88 | if array[0] is A { 89 | return ( 90 | array: array as! [A], 91 | shape: shape 92 | ) 93 | 94 | } else { 95 | let array = array as! [[Any]] 96 | shape.append(array[0].count) 97 | return flattenArrays( 98 | array: array.flatMap { $0 }, 99 | shape: &shape 100 | ) 101 | } 102 | } 103 | 104 | @inlinable 105 | public func splitRanges(total: Int, splits: Int) -> [Range] { 106 | let total = Float(total) 107 | let points = Set( 108 | (0 ... splits).map { (s: Int) -> Int in 109 | let percent = Float(s) / Float(splits) 110 | return Int(total * percent) 111 | } 112 | ) 113 | 114 | let orderedPoints = Array(points).sorted() 115 | var ranges = [Range]() 116 | 117 | for i in 0 ..< orderedPoints.count - 1 { 118 | ranges.append(orderedPoints[i] ..< orderedPoints[i + 1]) 119 | } 120 | 121 | return ranges 122 | } 123 | 124 | @inlinable 125 | public func getDimensionStrides(of shape: [Int]) -> [Int] { 126 | shape.reversed().scan(*).reversed().dropFirst() + [1] 127 | } 128 | 129 | @inlinable 130 | public func parFor( 131 | _ range: Range, 132 | workers: Int = CPU_COUNT, 133 | rangeMap: @escaping (Range) -> S, 134 | body: @escaping (E) -> Void 135 | ) where S.Element == E { 136 | let group = DispatchGroup() 137 | let nElements = range.count 138 | 139 | for range in splitRanges(total: nElements, splits: workers) { 140 | group.enter() 141 | 142 | DISPATCH.async { 143 | for element in rangeMap(range) { 144 | body(element) 145 | } 146 | group.leave() 147 | } 148 | } 149 | 150 | group.wait() 151 | } 152 | 153 | @inlinable 154 | public func parFor( 155 | _ range: Range, 156 | workers: Int = CPU_COUNT, 157 | body: @escaping (Int) -> Void 158 | ) { 159 | parFor(range, workers: workers, rangeMap: { x -> Range in x }, body: body) 160 | } 161 | 162 | extension Sequence { 163 | @inlinable 164 | func scan(initial: Element? = nil, _ f: @escaping (Element, Element) -> Element) -> AnySequence { 165 | AnySequence { () -> AnyIterator in 166 | var iterator = self.makeIterator() 167 | var acc = initial ?? iterator.next() 168 | var mayberNext = iterator.next() 169 | 170 | return AnyIterator { () -> Element? in 171 | if let next = mayberNext { 172 | defer { 173 | mayberNext = iterator.next() 174 | acc = f(acc!, next) 175 | } 176 | 177 | return acc 178 | } else { 179 | defer { 180 | acc = nil 181 | } 182 | return acc 183 | } 184 | } 185 | } 186 | } 187 | } 188 | 189 | extension Sequence where Element: Numeric { 190 | @inlinable 191 | public func sum() -> Element { 192 | reduce(0, +) 193 | } 194 | 195 | @inlinable 196 | public func product() -> Element { 197 | reduce(1, *) 198 | } 199 | } 200 | 201 | extension Sequence where Element == Bool { 202 | @inlinable 203 | public func all() -> Element { 204 | reduce(true) { $0 && $1 } 205 | } 206 | 207 | @inlinable 208 | public func product() -> Element { 209 | reduce(false) { $0 || $1 } 210 | } 211 | } -------------------------------------------------------------------------------- /Tests/LinuxMain.swift: -------------------------------------------------------------------------------- 1 | import XCTest 2 | 3 | import NDArrayTests 4 | 5 | var tests = [XCTestCaseEntry]() 6 | tests += NDArrayTests.allTests() 7 | XCTMain(tests) -------------------------------------------------------------------------------- /Tests/NDArrayTests/DimensionTests.swift: -------------------------------------------------------------------------------- 1 | // import Foundation 2 | @testable import NDArray 3 | import XCTest 4 | 5 | final class DimensionTests: XCTestCase { 6 | func testVirualEmpty() { 7 | let dimension = SingularDimension() 8 | 9 | let linearIndex = dimension.linearIndex(of: 0) 10 | let count = dimension.length 11 | 12 | XCTAssertEqual(linearIndex, 0) 13 | XCTAssertEqual(count, 1) 14 | } 15 | 16 | func testVirualOneRepeated() { 17 | let dimension = SingularDimension().tiled(3) 18 | 19 | let linearIndex = dimension.linearIndex(of: 1) 20 | let count = dimension.length 21 | let realCount = dimension.memory_layout.length 22 | 23 | XCTAssertEqual(linearIndex, 0) 24 | XCTAssertEqual(count, 3) 25 | XCTAssertEqual(realCount, 1) 26 | } 27 | 28 | func testVirualStrided() { 29 | let dimension = SingularDimension().tiled(10).sliced(stride: 3) 30 | 31 | let linearIndex = dimension.linearIndex(of: 1) 32 | let count = dimension.length 33 | let realCount = dimension.memory_layout.length 34 | 35 | XCTAssertEqual(linearIndex, 0) 36 | XCTAssertEqual(count, 4) 37 | XCTAssertEqual(realCount, 1) 38 | } 39 | 40 | func testVirualStridedMultiple() { 41 | let dimension = SingularDimension().tiled(4).sliced(stride: 3) 42 | 43 | let linearIndex = dimension.linearIndex(of: 1) 44 | let count = dimension.length 45 | let realCount = dimension.memory_layout.length 46 | 47 | XCTAssertEqual(linearIndex, 0) 48 | XCTAssertEqual(count, 2) 49 | XCTAssertEqual(realCount, 1) 50 | } 51 | 52 | func testVirualStridedMultiple2() { 53 | let dimension = SingularDimension().tiled(7).sliced(stride: 3) 54 | 55 | let linearIndex = dimension.linearIndex(of: 1) 56 | let count = dimension.length 57 | let realCount = dimension.memory_layout.length 58 | 59 | XCTAssertEqual(linearIndex, 0) 60 | XCTAssertEqual(count, 3) 61 | XCTAssertEqual(realCount, 1) 62 | } 63 | 64 | func testVirualStridedMultiple3() { 65 | let dimension = SingularDimension().tiled(8).sliced(stride: 3) 66 | 67 | let linearIndex = dimension.linearIndex(of: 1) 68 | let count = dimension.length 69 | let realCount = dimension.memory_layout.length 70 | 71 | XCTAssertEqual(linearIndex, 0) 72 | XCTAssertEqual(count, 3) 73 | XCTAssertEqual(realCount, 1) 74 | } 75 | 76 | func testVirualStridedMultiple4() { 77 | let dimension = SingularDimension().tiled(9).sliced(stride: 3) 78 | 79 | let linearIndex = dimension.linearIndex(of: 1) 80 | let count = dimension.length 81 | let realCount = dimension.memory_layout.length 82 | 83 | XCTAssertEqual(linearIndex, 0) 84 | XCTAssertEqual(count, 3) 85 | XCTAssertEqual(realCount, 1) 86 | } 87 | 88 | func testVirualStridedMultiple5() { 89 | let dimension = SingularDimension().tiled(9).sliced(start: 5, end: 6, stride: 3) 90 | 91 | let linearIndex = dimension.linearIndex(of: 0) 92 | let count = dimension.length 93 | let realCount = dimension.memory_layout.length 94 | 95 | XCTAssertEqual(linearIndex, 0) 96 | XCTAssertEqual(count, 1) 97 | XCTAssertEqual(realCount, 1) 98 | } 99 | 100 | func testVirualEmpty2() { 101 | let dimension = SingularDimension().tiled(3) 102 | 103 | let linearIndex = dimension.linearIndex(of: 1) 104 | 105 | XCTAssertEqual(linearIndex, 0) 106 | } 107 | 108 | func testScan() { 109 | let accumulatedProduct = [1, 2, 3, 4].scan(*) 110 | 111 | let total = accumulatedProduct.sum() 112 | 113 | XCTAssertEqual(total, 33) 114 | } 115 | 116 | func testShapeInit() { 117 | let shape = ArrayShape([5, 1, 4, 2]) 118 | 119 | XCTAssertEqual( 120 | shape.dimensions.map { $0.memory_layout.stride }, 121 | [8, 0, 2, 1] 122 | ) 123 | } 124 | 125 | static var allTests = [ 126 | ("testVirualEmpty", testVirualEmpty), 127 | ("testVirualOneRepeated", testVirualOneRepeated), 128 | ("testVirualStrided", testVirualStrided), 129 | ("testVirualStridedMultiple", testVirualStridedMultiple), 130 | ("testVirualStridedMultiple2", testVirualStridedMultiple2), 131 | ("testVirualStridedMultiple3", testVirualStridedMultiple3), 132 | ("testVirualStridedMultiple4", testVirualStridedMultiple4), 133 | ("testVirualStridedMultiple5", testVirualStridedMultiple5), 134 | ("testVirualEmpty2", testVirualEmpty2), 135 | ("testScan", testScan), 136 | ("testShapeInit", testShapeInit), 137 | ] 138 | } -------------------------------------------------------------------------------- /Tests/NDArrayTests/NDArrayTests.swift: -------------------------------------------------------------------------------- 1 | import Foundation 2 | @testable import NDArray 3 | import XCTest 4 | 5 | func timeIt(repetitions: Int = 1, function: () -> Void) -> Double { 6 | let startTime = Date() 7 | for _ in 1 ... repetitions { 8 | function() 9 | } 10 | return -startTime.timeIntervalSinceNow / Double(repetitions) 11 | } 12 | 13 | struct Point: AdditiveArithmetic { 14 | let x: Float 15 | let y: Float 16 | 17 | static var zero: Point { Point(x: 0, y: 0) } 18 | 19 | static prefix func + (lhs: Self) -> Self { 20 | lhs 21 | } 22 | 23 | static func + (lhs: Self, rhs: Self) -> Self { 24 | Point(x: lhs.x + rhs.x, y: lhs.y + rhs.y) 25 | } 26 | 27 | static func += (lhs: inout Self, rhs: Self) { 28 | lhs = lhs + rhs 29 | } 30 | 31 | static func - (lhs: Self, rhs: Self) -> Self { 32 | Point(x: lhs.x - rhs.x, y: lhs.y - rhs.y) 33 | } 34 | 35 | public static func -= (_ lhs: inout Point, _ rhs: Point) { 36 | lhs = lhs - rhs 37 | } 38 | } 39 | 40 | final class NDArrayTests: XCTestCase { 41 | func testElementWiseApply() { 42 | let a = NDArray([1, 2, 3], shape: [3]) 43 | let b = NDArray([1, 2, 3], shape: [3]) 44 | 45 | let c = a + b 46 | 47 | XCTAssert(c.toArray([Int].self) == [2, 4, 6]) 48 | } 49 | 50 | func testElementWiseApply2D() { 51 | let a = NDArray( 52 | [ 53 | 1, 2, 3, 54 | 4, 5, 6, 55 | ], 56 | shape: [2, 3] 57 | ) 58 | let b = NDArray( 59 | [ 60 | 1, 2, 3, 61 | 4, 5, 6, 62 | ], 63 | shape: [2, 3] 64 | ) 65 | 66 | let c = a + b 67 | 68 | XCTAssertEqual( 69 | c.toArray([[Int]].self), 70 | a.toArray([[Int]].self).map { arr in 71 | arr.map { x in x * 2 } 72 | } 73 | ) 74 | } 75 | 76 | func testConstructor() { 77 | let a = NDArray( 78 | [ 79 | [1, 2, 3], 80 | [4, 5, 6], 81 | ] 82 | ) 83 | let b = NDArray( 84 | [ 85 | [1, 2, 3], 86 | [4, 5, 6], 87 | ] 88 | ) 89 | let c = a + b 90 | 91 | XCTAssertEqual(a.shape, [2, 3]) 92 | XCTAssertEqual(b.shape, [2, 3]) 93 | XCTAssertEqual( 94 | c.toArray([[Int]].self), 95 | a.toArray([[Int]].self).map { arr in 96 | arr.map { x in x * 2 } 97 | } 98 | ) 99 | } 100 | 101 | func testElementWiseApply3D() { 102 | let a = NDArray( 103 | [ 104 | 1, 2, 3, 105 | 4, 5, 6, 106 | 107 | 7, 8, 9, 108 | 10, 11, 12, 109 | ], 110 | shape: [2, 2, 3] 111 | ) 112 | let b = NDArray( 113 | [ 114 | 1, 2, 3, 115 | 4, 5, 6, 116 | 117 | 7, 8, 9, 118 | 10, 11, 12, 119 | ], 120 | shape: [2, 2, 3] 121 | ) 122 | 123 | let c = a + b 124 | 125 | XCTAssertEqual( 126 | c.toArray([[[Int]]].self), 127 | a.toArray([[[Int]]].self).map { d0 in 128 | d0.map { d1 in 129 | d1.map { d2 in d2 * 2 } 130 | } 131 | } 132 | ) 133 | } 134 | 135 | func testElementWiseApply3DConstructor() { 136 | let a = NDArray( 137 | [ 138 | [ 139 | [1, 2, 3], 140 | [4, 5, 6], 141 | ], 142 | [ 143 | [7, 8, 9], 144 | [10, 11, 12], 145 | ], 146 | ] 147 | ) 148 | let b = NDArray( 149 | [ 150 | [ 151 | [1, 2, 3], 152 | [4, 5, 6], 153 | ], 154 | [ 155 | [7, 8, 9], 156 | [10, 11, 12], 157 | ], 158 | ] 159 | ) 160 | 161 | let c = a + b 162 | 163 | XCTAssertEqual(a.shape, [2, 2, 3]) 164 | XCTAssertEqual(b.shape, [2, 2, 3]) 165 | XCTAssertEqual(c.shape, [2, 2, 3]) 166 | XCTAssertEqual( 167 | c.toArray([[[Int]]].self), 168 | a.toArray([[[Int]]].self).map { d0 in 169 | d0.map { d1 in 170 | d1.map { d2 in d2 * 2 } 171 | } 172 | } 173 | ) 174 | } 175 | 176 | func testExample() { 177 | let a = NDArray( 178 | [ 179 | [1, 2, 3], 180 | [4, 5, 6], 181 | ] 182 | ) 183 | let b = NDArray( 184 | [ 185 | [7, 8, 9], 186 | [10, 11, 12], 187 | ] 188 | ) 189 | 190 | _ = (a + b) * a 191 | } 192 | 193 | func testCustomType() { 194 | let a = NDArray( 195 | [Point(x: 1, y: 2), Point(x: 2, y: 3)] 196 | ) 197 | let b = NDArray( 198 | [Point(x: 4, y: 5), Point(x: 6, y: 7)] 199 | ) 200 | 201 | let c = a + b 202 | 203 | let target = NDArray( 204 | [Point(x: 5, y: 7), Point(x: 8, y: 10)] 205 | ) 206 | 207 | XCTAssertEqual( 208 | c.toArray([Point].self), 209 | target.toArray([Point].self) 210 | ) 211 | } 212 | 213 | func testCustomType2() { 214 | let a = NDArray( 215 | [Point(x: 1, y: 2), Point(x: 2, y: 3)] 216 | ) 217 | let b = NDArray( 218 | [Point(x: 4, y: 5), Point(x: 6, y: 7)] 219 | ) 220 | 221 | let c: NDArray = elementwise(a, b, apply: +) 222 | 223 | let target = NDArray( 224 | [Point(x: 5, y: 7), Point(x: 8, y: 10)] 225 | ) 226 | 227 | XCTAssertEqual( 228 | c.toArray([Point].self), 229 | target.toArray([Point].self) 230 | ) 231 | } 232 | 233 | func testElementWiseApplyParallel() { 234 | let a = NDArray(Array(1 ... 100), shape: [100]) 235 | let b = NDArray(Array(1 ... 100), shape: [100]) 236 | 237 | let c = a + b 238 | 239 | XCTAssert( 240 | c.toArray([Int].self) == a.toArray([Int].self).map { $0 * 2 } 241 | ) 242 | } 243 | 244 | func testIndex() { 245 | let a = NDArray([ 246 | 3, 30, 247 | 2, 20, 248 | 1, 10, 249 | 250 | ], shape: [3, 2]) 251 | 252 | let b = a[1, 1] 253 | 254 | XCTAssertEqual(b.shape, []) 255 | XCTAssertEqual(b.scalarized(), 20) 256 | } 257 | 258 | func testScalarElementWiseAdd() { 259 | let a = NDArray([ 260 | 3, 30, 261 | 2, 20, 262 | 1, 10, 263 | 264 | ], shape: [3, 2]) 265 | 266 | let b = a[1, 1] 267 | let c = a[2, 0] 268 | 269 | let d = b + c 270 | 271 | XCTAssertEqual(b.baseCopy().data.value, [20]) 272 | XCTAssertEqual(c.baseCopy().data.value, [1]) 273 | XCTAssertEqual(d.toArray([Int].self), [21]) 274 | } 275 | 276 | func testBroadcast1() { 277 | let a = NDArray([1, 2, 3, 4], shape: [1, 4]) 278 | let b = NDArray([1, 2, 3, 4], shape: [4, 1]) 279 | 280 | let c = a + b 281 | 282 | let target = NDArray([ 283 | [2, 3, 4, 5], 284 | [3, 4, 5, 6], 285 | [4, 5, 6, 7], 286 | [5, 6, 7, 8], 287 | ]) 288 | 289 | XCTAssertEqual(c.baseCopy().data.value, target.baseCopy().data.value) 290 | } 291 | 292 | func testBroadcast2() { 293 | let a = NDArray([1, 2, 3, 4], shape: [1, 4]) 294 | let b = 1 295 | 296 | let c = a + b 297 | 298 | let target = NDArray([ 299 | 2, 3, 4, 5, 300 | ]) 301 | 302 | XCTAssertEqual(c.baseCopy().data.value, target.baseCopy().data.value) 303 | } 304 | 305 | func testBroadcast3() { 306 | let a = NDArray([1, 2, 3, 4], shape: [1, 4]) 307 | let b = 2 308 | 309 | let c = a * b 310 | 311 | let target = NDArray([ 312 | 2, 4, 6, 8, 313 | ]) 314 | 315 | XCTAssertEqual(c.baseCopy().data.value, target.baseCopy().data.value) 316 | } 317 | 318 | func testAssign() { 319 | var a = NDArray([1, 2, 3, 4], shape: [4]) 320 | var b = a 321 | 322 | a[0..] = NDArray([1, 1, 1, 1]) 323 | b[0..] = NDArray([2, 2, 2, 2]) 324 | 325 | XCTAssertEqual(a.baseCopy().data.value, [1, 1, 1, 1]) 326 | XCTAssertEqual(b.baseCopy().data.value, [2, 2, 2, 2]) 327 | } 328 | 329 | func testAssign2() { 330 | var a = NDArray([1, 2, 3, 4], shape: [4]) 331 | var b = a 332 | 333 | a[0..] = NDArray(1) 334 | b[0..] = NDArray(2) 335 | 336 | XCTAssertEqual(a.baseCopy().data.value, [1, 1, 1, 1]) 337 | XCTAssertEqual(b.baseCopy().data.value, [2, 2, 2, 2]) 338 | } 339 | 340 | func testAssign3() { 341 | var a = NDArray([1, 2, 3, 4], shape: [4]) 342 | var b = a 343 | 344 | a[0..] = NDArray(1) 345 | b[0..] = NDArray(2) 346 | 347 | XCTAssertEqual(a.baseCopy().data.value, [1, 1, 1, 1]) 348 | XCTAssertEqual(b.baseCopy().data.value, [2, 2, 2, 2]) 349 | } 350 | 351 | func testAssign4() { 352 | var a = NDArray([1, 2, 3, 4], shape: [4]) 353 | var b = a 354 | 355 | a[0..] = [1, 1, 1, 1] 356 | b[0..] = [2, 2, 2, 2] 357 | 358 | XCTAssertEqual(a.baseCopy().data.value, [1, 1, 1, 1]) 359 | XCTAssertEqual(b.baseCopy().data.value, [2, 2, 2, 2]) 360 | } 361 | 362 | func testAssign5() { 363 | var a = NDArray([1, 2, 3, 4], shape: [4]) 364 | var b = a 365 | var c = b[1..-1] 366 | 367 | a[0..] = [1, 1, 1, 1] 368 | b[0..] = [2, 2, 2, 2] 369 | 370 | let d = c[0..] 371 | 372 | c[0..] = NDArray(10) 373 | 374 | XCTAssertEqual(a.baseCopy().data.value, [1, 1, 1, 1]) 375 | XCTAssertEqual(b.baseCopy().data.value, [2, 2, 2, 2]) 376 | XCTAssertEqual(c.baseCopy().data.value, [10, 10]) 377 | XCTAssertEqual(d.baseCopy().data.value, [2, 3]) 378 | } 379 | 380 | func testTransposed() { 381 | let a = NDArray([ 382 | [1, 2, 3], 383 | [4, 5, 6], 384 | ]).transposed([1, 0]) 385 | 386 | XCTAssertEqual(a.baseCopy().data.value, [1, 4, 2, 5, 3, 6]) 387 | } 388 | 389 | func testRangesSplit() { 390 | let ranges = splitRanges(total: 70, splits: 11) 391 | 392 | XCTAssert(ranges.count == 11) 393 | } 394 | 395 | func testSqueezeAxis() { 396 | let a = NDArray([ 397 | [1, 2, 3, 4], 398 | ]) 399 | 400 | let b = a[squeeze, all] 401 | 402 | XCTAssertEqual(a.shape, [1, 4]) 403 | XCTAssertEqual(b.shape, [4]) 404 | } 405 | 406 | func testNewAxis() { 407 | let a = NDArray( 408 | [1, 2, 3, 4] 409 | ) 410 | let b = a[new] 411 | 412 | XCTAssertEqual(a.shape, [4]) 413 | XCTAssertEqual(b.shape, [1, 4]) 414 | } 415 | 416 | func testAll() { 417 | let a = NDArray([ 418 | [1, 2, 3, 4], 419 | [2, 3, 4, 5], 420 | ]) 421 | let b = a[ArrayRange.all, 1] 422 | 423 | XCTAssertEqual(b.baseCopy().data.value, [2, 3]) 424 | } 425 | 426 | func testEllipsis() { 427 | let a = NDArray( 428 | Array(1 ... 16), 429 | shape: [2, 2, 2, 2] 430 | ) 431 | let b = a[ArrayRange.ellipsis] 432 | 433 | XCTAssertEqual(b.shape, [2, 2, 2, 2]) 434 | } 435 | 436 | func testEllipsis2() { 437 | let a = NDArray( 438 | Array(1 ... 16), 439 | shape: [2, 2, 2, 2] 440 | ) 441 | let b = a[0, ArrayRange.ellipsis] 442 | 443 | XCTAssertEqual(b.shape, [2, 2, 2]) 444 | } 445 | 446 | func testEllipsis3() { 447 | let a = NDArray( 448 | Array(1 ... 16), 449 | shape: [2, 2, 2, 2] 450 | ) 451 | let b = a[0, rest, 0] 452 | 453 | XCTAssertEqual(b.shape, [2, 2]) 454 | } 455 | 456 | func testEllipsis4() { 457 | let a = NDArray( 458 | Array(1 ... 16), 459 | shape: [2, 2, 2, 2] 460 | ) 461 | let b = a[0, ArrayRange.ellipsis, 0, 0] 462 | 463 | XCTAssertEqual(b.shape, [2]) 464 | } 465 | 466 | func testEllipsis5() { 467 | let a = NDArray( 468 | Array(1 ... 16), 469 | shape: [2, 2, 2, 2] 470 | ) 471 | let b = a[0, 0, rest, 0, 0] 472 | 473 | XCTAssertEqual(b.shape, []) 474 | } 475 | 476 | func testNegativeStride() { 477 | let a = NDArray([1, 2, 3, 4, 5]) 478 | 479 | let b = a[((-1)...).stride(-1)] 480 | 481 | XCTAssertEqual(a.baseCopy().data.value, b.baseCopy().data.value.reversed()) 482 | } 483 | 484 | func testNegativeStride2() { 485 | let a = NDArray([1, 2, 3, 4, 5]) 486 | 487 | let b = a[....-1] 488 | 489 | XCTAssertEqual(a.baseCopy().data.value, b.baseCopy().data.value.reversed()) 490 | } 491 | 492 | func testNegativeStride3() { 493 | let a = NDArray([1, 2, 3, 4, 5]) 494 | 495 | let b = a[..1..-2] 496 | 497 | XCTAssertEqual(b.baseCopy().data.value, [5, 3]) 498 | } 499 | 500 | func testNegativeIndex() { 501 | let a = NDArray([1, 2, 3, 4, 5]) 502 | 503 | let b = a[-1] 504 | 505 | XCTAssertEqual(b.baseCopy().data.value, [5]) 506 | } 507 | 508 | func testFilter() { 509 | let a = NDArray([1, 2, 3, 4, 5]) 510 | 511 | let b = a[[1, 1, 0, 2, -1]] 512 | 513 | XCTAssertEqual(b.baseCopy().data.value, [2, 2, 1, 3, 5]) 514 | } 515 | 516 | func testFilter2() { 517 | let a = NDArray([1, 2, 3, 4, 5]) 518 | 519 | let b = a[[true, true, false, false, true]] 520 | 521 | XCTAssertEqual(b.baseCopy().data.value, [1, 2, 5]) 522 | } 523 | 524 | func testReduce1() { 525 | let a = NDArray([ 526 | [1, 2, 3, 4], 527 | [10, 20, 30, 40], 528 | ]) 529 | 530 | let b = a.reduce( 531 | axis: [0], 532 | initial: NDArray([0, 0, 0, 0]), 533 | f: + 534 | ) 535 | 536 | XCTAssertEqual(b.shape, [4]) 537 | XCTAssertEqual(b.baseCopy().data.value, [11, 22, 33, 44]) 538 | } 539 | 540 | func testReduce2() { 541 | let a = NDArray([ 542 | [1, 2, 3, 4], 543 | [10, 20, 30, 40], 544 | ]) 545 | 546 | let b = a.reduce( 547 | axis: [1], 548 | initial: NDArray([0, 0]), 549 | f: + 550 | ) 551 | 552 | XCTAssertEqual(b.shape, [2]) 553 | XCTAssertEqual(b.baseCopy().data.value, [10, 100]) 554 | } 555 | 556 | func testReduce3() { 557 | let a = NDArray([ 558 | [1, 2, 3, 4], 559 | [10, 20, 30, 40], 560 | ]) 561 | 562 | let b = a.reduce( 563 | axis: [0, 1], 564 | initial: NDArray(0), 565 | f: + 566 | ) 567 | 568 | XCTAssertEqual(b.shape, []) 569 | XCTAssertEqual(b.scalarized(), 110) 570 | } 571 | 572 | func testSum1() { 573 | let a = NDArray([ 574 | [1, 2, 3, 4], 575 | [10, 20, 30, 40], 576 | ]) 577 | 578 | let b = a.sum( 579 | axis: [0] 580 | ) 581 | 582 | XCTAssertEqual(b.shape, [4]) 583 | XCTAssertEqual(b.baseCopy().data.value, [11, 22, 33, 44]) 584 | } 585 | 586 | func testSum2() { 587 | let a = NDArray([ 588 | [1, 2, 3, 4], 589 | [10, 20, 30, 40], 590 | ]) 591 | 592 | let b = a.sum( 593 | axis: [1] 594 | ) 595 | 596 | XCTAssertEqual(b.shape, [2]) 597 | XCTAssertEqual(b.baseCopy().data.value, [10, 100]) 598 | } 599 | 600 | func testSum3() { 601 | let a = NDArray([ 602 | [1, 2, 3, 4], 603 | [10, 20, 30, 40], 604 | ]) 605 | 606 | let b = a.sum( 607 | axis: [0, 1] 608 | ) 609 | 610 | XCTAssertEqual(b.shape, []) 611 | XCTAssertEqual(b.scalarized(), 110) 612 | } 613 | 614 | func testMean1() { 615 | let a: NDArray = [ 616 | [1, 2, 3, 4], 617 | [10, 20, 30, 40], 618 | ] 619 | 620 | let b = a.mean( 621 | axis: [0] 622 | ) 623 | 624 | XCTAssertEqual(b.shape, [4]) 625 | XCTAssertEqual(b.baseCopy().data.value, [11 / 2, 22 / 2, 33 / 2, 44 / 2] as [Float]) 626 | } 627 | 628 | func testMean2() { 629 | let a: NDArray = [ 630 | [1, 2, 3, 4], 631 | [10, 20, 30, 40], 632 | ] 633 | 634 | let b = a.mean( 635 | axis: [1] 636 | ) 637 | 638 | XCTAssertEqual(b.shape, [2]) 639 | XCTAssertEqual(b.baseCopy().data.value, [10 / 4, 100 / 4] as [Float]) 640 | } 641 | 642 | func testMean3() { 643 | let a = NDArray([[Float]]([ 644 | [1, 2, 3, 4], 645 | [10, 20, 30, 40], 646 | ])) 647 | 648 | let b = a.mean( 649 | axis: [0, 1] 650 | ) 651 | 652 | XCTAssertEqual(b.shape, []) 653 | XCTAssertEqual(b.scalarized(), 110 / 8) 654 | } 655 | 656 | func testMax1() { 657 | let a = NDArray([ 658 | [1, 2, 3, 4], 659 | [10, 20, 30, 40], 660 | ]) 661 | 662 | let b = a.max( 663 | axis: [0] 664 | ) 665 | 666 | XCTAssertEqual(b.shape, [4]) 667 | XCTAssertEqual(b.baseCopy().data.value, [10, 20, 30, 40]) 668 | } 669 | 670 | func testMax2() { 671 | let a = NDArray([ 672 | [1, 2, 3, 4], 673 | [10, 20, 30, 40], 674 | ]) 675 | 676 | let b = a.max( 677 | axis: [1] 678 | ) 679 | 680 | XCTAssertEqual(b.shape, [2]) 681 | XCTAssertEqual(b.baseCopy().data.value, [4, 40]) 682 | } 683 | 684 | func testMax3() { 685 | let a = NDArray([ 686 | [1, 2, 3, 4], 687 | [10, 20, 30, 40], 688 | ]) 689 | 690 | let b = a.max( 691 | axis: [0, 1] 692 | ) 693 | 694 | XCTAssertEqual(b.shape, []) 695 | XCTAssertEqual(b.scalarized(), 40) 696 | } 697 | 698 | func testMin1() { 699 | let a = NDArray([ 700 | [1, 2, 3, 4], 701 | [10, 20, 30, 40], 702 | ]) 703 | 704 | let b = a.min( 705 | axis: [0] 706 | ) 707 | 708 | XCTAssertEqual(b.shape, [4]) 709 | XCTAssertEqual(b.baseCopy().data.value, [1, 2, 3, 4]) 710 | } 711 | 712 | func testMin2() { 713 | let a = NDArray([ 714 | [1, 2, 3, 4], 715 | [10, 20, 30, 40], 716 | ]) 717 | 718 | let b = a.min( 719 | axis: [1] 720 | ) 721 | 722 | XCTAssertEqual(b.shape, [2]) 723 | XCTAssertEqual(b.baseCopy().data.value, [1, 10]) 724 | } 725 | 726 | func testMin3() { 727 | let a = NDArray([ 728 | [1, 2, 3, 4], 729 | [10, 20, 30, 40], 730 | ]) 731 | 732 | let b = a.min( 733 | axis: [0, 1] 734 | ) 735 | 736 | XCTAssertEqual(b.shape, []) 737 | XCTAssertEqual(b.scalarized(), 1) 738 | } 739 | 740 | static var allTests = [ 741 | ("testElementWiseApply", testElementWiseApply), 742 | ("testElementWiseApply2D", testElementWiseApply2D), 743 | ("testElementWiseApply3D", testElementWiseApply3D), 744 | ("testElementWiseApplyParallel", testElementWiseApplyParallel), 745 | // ("testElementWiseApplyParallelBenchmark", testElementWiseApplyParallelBenchmark), 746 | ("testRangesSplit", testRangesSplit), 747 | ("testIndex", testIndex), 748 | ("testScalarElementWiseAdd", testScalarElementWiseAdd), 749 | ("testConstructor", testConstructor), 750 | ("testElementWiseApply3DConstructor", testElementWiseApply3DConstructor), 751 | ("testTransposed", testTransposed), 752 | ("testExample", testExample), 753 | ("testCustomType", testCustomType), 754 | ("testCustomType2", testCustomType2), 755 | ("testBroadcast1", testBroadcast1), 756 | ("testBroadcast2", testBroadcast2), 757 | ("testBroadcast3", testBroadcast3), 758 | ("testAssign", testAssign), 759 | ("testAssign2", testAssign2), 760 | ("testAssign3", testAssign3), 761 | ("testAssign4", testAssign4), 762 | ("testAssign5", testAssign5), 763 | ("testSqueezeAxis", testSqueezeAxis), 764 | ("testNewAxis", testNewAxis), 765 | ("testAll", testAll), 766 | ("testEllipsis", testEllipsis), 767 | ("testEllipsis2", testEllipsis2), 768 | ("testEllipsis3", testEllipsis3), 769 | ("testEllipsis4", testEllipsis4), 770 | ("testEllipsis5", testEllipsis5), 771 | ("testNegativeStride", testNegativeStride), 772 | ("testNegativeStride2", testNegativeStride2), 773 | ("testNegativeStride3", testNegativeStride3), 774 | ("testNegativeIndex", testNegativeIndex), 775 | ("testFilter", testFilter), 776 | ("testFilter2", testFilter2), 777 | ("testReduce1", testReduce1), 778 | ("testReduce2", testReduce2), 779 | ("testReduce3", testReduce3), 780 | ("testSum1", testSum1), 781 | ("testSum2", testSum2), 782 | ("testSum3", testSum3), 783 | ("testMean1", testMean1), 784 | ("testMean2", testMean2), 785 | ("testMean3", testMean3), 786 | ("testMax1", testMax1), 787 | ("testMax2", testMax2), 788 | ("testMax3", testMax3), 789 | ("testMin1", testMin1), 790 | ("testMin2", testMin2), 791 | ("testMin3", testMin3), 792 | ] 793 | } -------------------------------------------------------------------------------- /Tests/NDArrayTests/XCTestManifests.swift: -------------------------------------------------------------------------------- 1 | import XCTest 2 | 3 | #if !canImport(ObjectiveC) 4 | public func allTests() -> [XCTestCaseEntry] { 5 | return [ 6 | testCase(NDArrayTests.allTests), 7 | testCase(DimensionTests.allTests), 8 | ] 9 | } 10 | #endif --------------------------------------------------------------------------------