Swift Tensor API Proposal

tensor-api.swift

//
//  CeedNumerics
//
//  Created by Raphael Sebbe on 16/11/2018.
//  Copyright © 2018 Creaceed. All rights reserved.
//

import Foundation

// Slicing

// Unresolved one (needs shape to resolve). In NumPy,
// <start:end:step> (end is excluded)
// <2:> -> means <2 to end>.
// <:> -> means <all>.
// <:3> -> means <start to 3>.
// <3:0:-1> -> means <3,2,1>.
// <3::-1> -> means <3,2,1,0>.
// <0:6:2> -> means <0,2,4>.
// start, end, step can all be negative.
//
// Serves as base type for subscript slicing.
public protocol NSliceExpression {
    public var start: Int? { get }
    public var end: Int? { get }
    public var step: Int? { get }
}

public extension NSliceExpression {
    public func resolve(size: Int) -> NResolvedSlice
    public func resolve(within parent: NResolvedSlice) -> NResolvedSlice
}

public struct NSlice : NSliceExpression {
    public let start: Int?
    public let end: Int?
    public let step: Int?
    public static let all: NSlice
  
    public init(start: Int?, end: Int?, step: Int?)
}

public struct NResolvedSlice : NSliceExpression {
    public let rstart: Int
    public let rcount: Int
    public let rstep: Int
    public var rlast: Int { get }
    public var start: Int? { get }
    public var end: Int? { get }
    public var step: Int? { get }

    public init(start: Int, count: Int, step: Int)
    public static func `default`(count: Int) -> NResolvedSlice
    public func position(at index: Int) -> Int
    public func flatten(within parent: NResolvedSlice) -> NResolvedSlice
}

// Storage (Memory)
public class NStorage<Element> where Element : NValue {
    public let count: Int
    public struct Access {
        public var base: UnsafeMutablePointer<Element>
        public var count: Int
    }

    public init(allocatedCount: Int)
    public init(externalReference: UnsafeMutablePointer<Element>, count c: Int)
    public subscript(index: Int) -> Element { get set }
    public func withUnsafeAccess<Result>(_ block: (Access) throws -> Result) rethrows -> Result
}

// Dimension Type base protocol (for vector, matrix, tensor)
public protocol NDimensionalType : CustomStringConvertible {
    associatedtype Element : NValue
    public var dimension: Int { get }
    public var shape: [Int] { get }
    public subscript(index: [Int]) -> Element { get set }
    public func isCompact(dimension: Int) -> Bool
}

extension NDimensionalType {
    public var description: String { get }
}

public class DimensionalIterator : IteratorProtocol {
    public init(shape: [Int])
    public func next() -> [Int]?
}

// 1-D Vector
public struct NVectorLayout {
    public let offset: Int
    public let stride: Int
    public static let `default`: CeedNumerics.NVectorLayout
  
    public func location(for position: Int) -> Int
}

public class NVector<Element> where Element : NValue {
    public typealias Storage = NStorage<Element>
  
    public var size: Int { get }

    public init(storage mem: Storage, layout l: NVectorLayout = default, slice sl: NResolvedSlice)
    public convenience init(size: Int)
    public convenience init(storage mem: Storage, layout l: NVectorLayout = default, count: Int)
    public subscript(slice s: NSliceExpression) -> NVector { get }
    public subscript(index: Int) -> Element { get set }
}

extension NVector where Element : NValue {
    public func set(from: NVector)
}

extension NVector : NDimensionalType where Element : NValue {
    public var dimension: Int { get }
    public var shape: [Int] { get }
  
    public subscript(index: [Int]) -> Element { get set }
    public func isCompact(dimension: Int) -> Bool
}

// 2-D Matrix
public struct NMatrixLayout {
    public let offset: Int
    public let stride: (row: Int, column: Int)
    public static func `default`(rows: Int) -> NMatrixLayout
  
    public func location(row rpos: Int, column cpos: Int) -> Int
    public func layout(row: Int) -> NVectorLayout
    public func layout(column: Int) -> NVectorLayout
}

public class NMatrix<Element> where Element : NValue {
    public typealias Storage = NStorage<Element>
    public typealias Vector = NVector<Element>
    public typealias Matrix = NMatrix<Element>
    public var rows: Int { get }
    public var columns: Int { get }

    public init(storage s: Storage, layout l: NMatrixLayout, slices sl: (NResolvedSlice, NResolvedSlice))
    public convenience init(rows: Int, columns: Int)
    public subscript(row row: Int) -> Vector { get set }
    public subscript(column col: Int) -> Vector { get }
    public subscript(rowSlice: NSliceExpression, colSlice: NSliceExpression) -> Matrix { get }
    public subscript(row: Int, column: Int) -> Element { get set }
}

extension NMatrix : NDimensionalType where Element : NValue {
    public var dimension: Int { get }
    public var shape: [Int] { get }
  
    public subscript(index: [Int]) -> Element { get set }
    public func isCompact(dimension: Int) -> Bool
}


// N-D Tensor --- not yet done.