8d105685创建于 2025年10月30日历史提交
package scientific.linear

import std.math.*
import std.time.*
import std.collection.*
import std.sort.*
import std.unittest.*
import std.unittest.testmacro.*

import scientific.numbers.*
import scientific.utils.assertEqual
import scientific.utils.AssertionException

foreign func malloc(size: UIntNative): CPointer<Unit>
foreign func free(ptr: CPointer<Unit>): Unit

public abstract class AbstractTensor<T> where T <: CNumber<T> {
    /*
     * List of dimensions of the tensor.
     */
    var dims: Array<Int64> = Array<Int64>()

    /*
     * Obtain the size of the tensor (in number of elements).
     */
    public func size(): Int64 {
        var res = 1
        for (i in this.dims) {
            res = res * i
        }
        return res
    }

    /*
     * Obtain the dimensions of the tensor
     */
    public func shape(): Array<Int64> {
        return this.dims
    }

    /*
     * Return a string representing type of elements.
     */
    public func dtype(): String {
        return T.getType()
    }
}

public class Vector<T> <: AbstractTensor<T> & Equatable<Vector<T>> where T <: CNumber<T> {
    /*
     * Address storing data.
     */
    public var ptr: CPointer<Unit> = CPointer<Unit>()

    /*
     * Initialization from a pointer.
     */
    init(ptr: CPointer<Unit>, size: UIntNative) {
        this.ptr = ptr
        this.dims = [Int64(size)]
    }

    /*
     * Constructor from length of vector.
     */
    init(size: Int64) {
        if (size < 0) {
            throw IllegalArgumentException("init: size must be non-negative.")
        }
        this.dims = [size]
        if (size > 0) {
            this.ptr = unsafe { malloc(UIntNative(T.getSize() * size)) }
        }
    }

    /*
     * Constructor from length of vector and value.
     */
    init(size: Int64, x: T) {
        if (size < 0) {
            throw IllegalArgumentException("init: size must be non-negative.")
        }
        this.dims = [size]
        if (size > 0) {
            this.ptr = unsafe { malloc(UIntNative(T.getSize() * size)) }
            for (i in 0..size) {
                this[i] = x
            }
        }
    }

    /*
     * Constructor from length of vector and initialization function.
     */
    init(size: Int64, f: (Int64) -> T) {
        if (size < 0) {
            throw IllegalArgumentException("init: size must be non-negative.")
        }
        this.dims = [size]
        if (size > 0) {
            this.ptr = unsafe { malloc(UIntNative(T.getSize() * size)) }
            for (i in 0..size) {
                this[i] = f(i)
            }
        }
    }

    /*
     * Constructor from an Array.
     */
    init(a: Array<T>) {
        var size = a.size
        this.dims = [size]
        if (size > 0) {
            this.ptr = unsafe { malloc(UIntNative(T.getSize() * size)) }
            for (i in 0..size) {
                this[i] = a[i]
            }
        }
    }

    /*
     * Make a copy of the vector.
     */
    public func copy(): Vector<T> {
        let size = this.size()
        var res = Vector<T>(size)
        for (i in 0..size) {
            res[i] = this[i]
        }
        return res
    }

    /*
     * Access i'th element of a vector.
     */
    public operator func [](index: Int64): T {
        if (0 <= index && index < this.size()) {
            return unsafe { T.read(this.ptr, index) }
        } else {
            throw IllegalArgumentException("Index out of range")
        }
    }

    /*
     * Write i'th element of the vector.
     */
    public operator func [](index: Int64, value!: T): Unit {
        if (0 <= index && index < this.size()) {
            unsafe { T.write(this.ptr, value, index) }
        } else {
            throw IllegalArgumentException("Index out of range")
        }
    }

    /*
     * Access range of a vector.
     */
    public operator func [](range: Range<Int64>): Vector<T> {
        var a = ArrayList<T>()
        var size = this.size()
        for (i in range) {
            if (i >= size) {
                break
            }
            a.add(this[i])
        }
        return Vector<T>(a.toArray())
    }

    /*
     * Apply function f to every element of the vector.
     */
    public func apply(f: (T) -> T): Vector<T> {
        var size = this.size()
        var res = Vector<T>(size)
        for (i in 0..size) {
            res[i] = f(this[i])
        }
        return res
    }

    /*
     * Filter by condition.
     */
    public func filter(f: (T) -> Bool): Vector<T> {
        var a = ArrayList<T>()
        var size = this.size()
        for (i in 0..size) {
            if (f(this[i])) {
                a.add(this[i])
            }
        }
        return Vector<T>(a.toArray())
    }

    public operator func ==(a: Vector<T>): Bool {
        var size = this.size()
        if (size != a.size()) {
            return false
        }
        for (i in 0..size) {
            if (this[i] != a[i]) {
                return false
            }
        }
        return true
    }

    public operator func !=(a: Vector<T>): Bool {
        return !(this == a)
    }

    /*
     * Add two vectors element-wise.
     */
    public operator func +(a: Vector<T>): Vector<T> {
        var size = this.size()
        if (size != a.size()) {
            throw IllegalArgumentException("plus: size mismatch")
        }
        var res = Vector<T>(size)
        for (i in 0..size) {
            res[i] = this[i] + a[i]
        }
        return res
    }

    /*
     * Add scalar to vector.
     */
    public operator func +(a: T): Vector<T> {
        return this.apply({x => x + a})
    }
    
    /*
     * Subtract two vectors element-wise.
     */
    public operator func -(a: Vector<T>): Vector<T> {
        var size = this.size()
        if (size != a.size()) {
            throw IllegalArgumentException("minus: size mismatch")
        }
        var res = Vector<T>(size)
        for (i in 0..size) {
            res[i] = this[i] - a[i]
        }
        return res
    }

    /*
     * Subtract scalar from vector.
     */
    public operator func -(a: T): Vector<T> {
        return this.apply({x => x - a})
    }

    /*
     * Multiply two vectors element-wise.
     */
    public operator func *(a: Vector<T>): Vector<T> {
        var size = this.size()
        if (size != a.size()) {
            throw IllegalArgumentException("multiply: size mismatch")
        }
        var res = Vector<T>(size)
        for (i in 0..size) {
            res[i] = this[i] * a[i]
        }
        return res
    }

    /*
     * Scalar multiplication with a constant.
     */
    public operator func *(a: T): Vector<T> {
        return this.apply({x => x * a})
    }

    /*
     * Divide a vector by a scalar.
     */
    public operator func /(a: T): Vector<T> {
        return this.apply({x => x / a})
    }

    /*
     * Divide a vector by a vector.
     */
    public operator func /(a: Vector<T>): Vector<T> {
        let size = this.size()
        if (size != a.size()) {
            throw IllegalArgumentException("divides: size do not match")
        }
        return Vector<T>(size, {i => this[i] / a[i]})
    }

    /*
     * Dot product.
     */
    public func dot(a: Vector<T>): T {
        var size = this.size()
        if (size != a.size()) {
            throw IllegalArgumentException("dot: size mismatch")
        }
        var res = T.fromInt(0)
        for (i in 0..size) {
            res += this[i] * a[i]
        }
        return res
    }

    public func toArray(): Array<T> {
        let size = this.size()
        return Array<T>(size, {i => this[i]})
    }

    /*
     * Convert vector of length n to an n-by-1 matrix.
     */
    public func toMatrix(): Matrix<T> {
        var size = this.size()
        var res = Matrix<T>(size, 1)
        for (i in 0..size) {
            res[i,0] = this[i]
        }
        return res
    }

    ~init() {
        if (!ptr.isNull()) {
            unsafe { free(ptr) }
        }
    }
}

public func empty<T>(size: Int64): Vector<T> where T <: CNumber<T> {
    return Vector<T>(size)
}

public func zeros<T>(size: Int64): Vector<T> where T <: CNumber<T> {
    return Vector<T>(size, T.fromInt(0))
}

public func ones<T>(size: Int64): Vector<T> where T <: CNumber<T> {
    return Vector<T>(size, T.fromInt(1))
}

public func vector<T>(size: Int64, x: T): Vector<T> where T <: CNumber<T> {
    return Vector<T>(size, x)
}

public func vector<T>(size: Int64, f: (Int64) -> T): Vector<T> where T <: CNumber<T> {
    return Vector<T>(size, f)
}

public func vector<T>(a: Array<T>): Vector<T> where T <: CNumber<T> {
    return Vector<T>(a)
}

public func vector<T>(ptr: CPointer<Unit>, size: Int64): Vector<T> where T <: CNumber<T> {
    return Vector<T>(ptr, UIntNative(size))
}

extend<T> Vector<T> <: ToString where T <: ToString {
    public func toString(): String {
        let builder = StringBuilder()
        builder.append("[")
        let size = this.size()
        for (i in 0..size) {
            if (i != 0) {
                builder.append(", ")
            }
            builder.append(this[i].toString())
        }
        builder.append("]")
        return builder.toString()
    }
}

extend<T> Vector<T> where T <: Float<T> {
    /*
     * Norm of a vector.
     */
    public func norm(): T {
        return sqrt(this.dot(this))
    }
}

extend<T> Vector<T> where T <: Real<T> {
    /*
     * Convert the vector to elements of type Float64.
     */
    public func toFloat64(): Vector<Float64> {
        let size = this.size()
        return Vector<Float64>(size, {i => this[i].toFloat64()})
    }
}

extend<T> Vector<T> where T <: Comparable<T> {
    /*
     * Clip content of a vector.
     */
    public func clip(min: T, max: T): Vector<T> {
        let size = this.size()
        let res = empty<T>(size)
        for (i in 0..size) {
            if (this[i] < min) {
                res[i] = min
            } else if (this[i] > max) {
                res[i] = max
            } else {
                res[i] = this[i]
            }
        }
        return res
    }
}

public func abs<T>(x: Vector<T>): Vector<T> where T <: Real<T> {
    let size = x.size()
    return Vector<T>(size, {i => abs(x[i])})
}

public func power<T>(base: Vector<T>, exp: T): Vector<T> where T <: Float<T> {
    let size = base.size()
    return Vector<T>(size, {i => power(base[i], exp)})
}

public class Matrix<T> <: AbstractTensor<T> & Equatable<Matrix<T>> where T <: CNumber<T> {
    /*
     * Address storing data.
     */
    public var ptr: CPointer<Unit> = CPointer<Unit>()

    /*
     * Constructor given number of rows and columns.
     */
    init(rows: Int64, cols: Int64) {
        if (rows < 0 || cols < 0) {
            throw IllegalArgumentException("init: rows and cols must be nonnegative.")
        }
        this.dims = [rows, cols]
        if (rows * cols > 0) {
            this.ptr = unsafe { malloc(UIntNative(T.getSize() * rows * cols)) }
        }
    }

    init(ptr: CPointer<Unit>, rows: UIntNative, cols: UIntNative) {
        this.ptr = ptr
        this.dims = [Int64(rows), Int64(cols)]
    }

    /*
     * Constructor from dimensions and initial value.
     */
    init(rows: Int64, cols: Int64, x: T) {
        if (rows < 0 || cols < 0) {
            throw IllegalArgumentException("init: rows and cols must be nonnegative.")
        }
        this.dims = [rows, cols]
        if (rows * cols > 0) {
            this.ptr = unsafe { malloc(UIntNative(T.getSize() * rows * cols)) }
        }
        for (i in 0..rows) {
            for (j in 0..cols) {
                this[i,j] = x
            }
        }
    }

    /*
     * Constructor from dimensions and initialization function.
     */
    init(rows: Int64, cols: Int64, f: (Int64, Int64) -> T) {
        if (rows < 0 || cols < 0) {
            throw IllegalArgumentException("init: rows and cols must be nonnegative.")
        }
        this.dims = [rows, cols]
        if (rows * cols > 0) {
            this.ptr = unsafe { malloc(UIntNative(T.getSize() * rows * cols)) }
        }
        for (i in 0..rows) {
            for (j in 0..cols) {
                this[i,j] = f(i, j)
            }
        }
    }

    /*
     * Constructor from an array of arrays.
     */
    init(a: Array<Array<T>>) {
        var rows = a.size
        if (rows == 0) {
            throw IllegalArgumentException("init: empty matrix.")
        }
        var cols = a[0].size
        for (i in 1..rows) {
            if (a[i].size != cols) {
                throw IllegalArgumentException(
                    "init: column size mismatch, ${cols} vs ${a[i].size}.")
            }
        }

        this.dims = [rows, cols]
        if (rows * cols > 0) {
            this.ptr = unsafe { malloc(UIntNative(T.getSize() * rows * cols)) }
        }
        for (i in 0..rows) {
            for (j in 0..cols) {
                this[i,j] = a[i][j]
            }
        }
    }

    /*
     * Obtain number of rows in a matrix.
     */
    public func getRows(): Int64 {
        return this.dims[0]
    }

    /*
     * Obtain number of columns in a matrix.
     */
    public func getCols(): Int64 {
        return this.dims[1]
    }

    /*
     * Whether the matrix is a square matrix.
     */
    public func isSquare(): Bool {
        return this.dims[0] == this.dims[1]
    }

    /*
     * Obtain a copy of the matrix.
     */
    public func copy(): Matrix<T> {
        var res = Matrix<T>(this.getRows(), this.getCols())
        let row = this.getRows()
        let col = this.getCols()
        for (i in 0..row) {
            for (j in 0..col) {
                res[i,j] = this[i,j]
            }
        }
        return res
    }

    /*
     * Access i,j'th element of a matrix.
     */
    public operator func [](r: Int64, c: Int64): T {
        if (0 <= r && r < this.getRows() && 0 <= c && c < this.getCols()) {
            return unsafe { T.read(this.ptr, r + c * this.dims[0]) }
        } else {
            throw IllegalArgumentException("Index out of range")
        }
    }

    /*
     * Write i,j'th element of a matrix.
     */
    public operator func [](r: Int64, c: Int64, value!: T): Unit {
        if (0 <= r && r < this.getRows() && 0 <= c && c < this.getCols()) {
            unsafe { T.write(this.ptr, value, r + c * this.dims[0]) }
        } else {
            throw IllegalArgumentException("Index out of range")
        }
    }

    /*
     * Read i'th row of a matrix.
     */
    public operator func [](i: Int64): Vector<T> {
        if (0 <= i && i < this.getRows()) {
            var res = Vector<T>(this.getCols())
            for (j in 0..this.getCols()) {
                res[j] = this[i,j]
            }
            return res
        } else {
            throw IllegalArgumentException("Index out of range")
        }
    }

    /*
     * Write i'th row of a matrix.
     */
    public operator func [](i: Int64, value!: Vector<T>): Unit {
        if (value.size() != this.getCols()) {
            throw IllegalArgumentException("dimension mismatch")
        }
        if (0 <= i && i < this.getRows()) {
            for (j in 0..this.getCols()) {
                this[i,j] = value[j]
            }
        } else {
            throw IllegalArgumentException("Index out of range")
        }
    }

    public func getCol(i: Int64): Vector<T> {
        if (0 <= i && i < this.getCols()) {
            var res = Vector<T>(this.getRows())
            for (j in 0..this.getRows()) {
                res[j] = this[j,i]
            }
            return res
        } else {
            throw IllegalArgumentException("Index out of range")
        }
    }

    public func setCol(i: Int64, value: Vector<T>): Unit {
        if (value.size() != this.getRows()) {
            throw IllegalArgumentException("dimension mismatch")
        }
        if (0 <= i && i < this.getCols()) {
            for (j in 0..this.getRows()) {
                this[j,i] = value[j]
            }
        } else {
            throw IllegalArgumentException("Index out of range")
        }
    }

    public operator func ==(a: Matrix<T>): Bool {
        if (this.shape() != a.shape()) {
            return false
        }
        let r = this.getRows()
        let c = this.getCols()
        for (i in 0..r) {
            for (j in 0..c) {
                if (this[i,j] != a[i,j]) {
                    return false
                }
            }
        }
        return true
    }

    public operator func !=(a: Matrix<T>): Bool {
        return !(this == a)
    }

    /*
     * Add two matrices element-wise.
     */
    public operator func +(a: Matrix<T>): Matrix<T> {
        if(this.shape() != a.shape()) {
            throw IllegalArgumentException("plus: dimension mismatch")
        }

        let r = this.getRows()
        let c = this.getCols()
        return Matrix<T>(r, c, {i, j => this[i,j] + a[i,j]})
    }

    /*
     * Subtract two matrices element-wise.
     */
    public operator func -(a: Matrix<T>): Matrix<T> {
        if(this.shape() != a.shape()) {
            throw IllegalArgumentException("plus: dimension mismatch")
        }

        let r = this.getRows()
        let c = this.getCols()
        return Matrix<T>(r, c, {i, j => this[i,j] - a[i,j]})
    }

    /*
     * Scalar multiplication of a matrix with a constant.
     */
    public operator func *(a: T): Matrix<T> {
        return this.apply({x => x * a})
    }

    /*
     * Apply function f to every element of the matrix.
     */
    public func apply(f: (T) -> T): Matrix<T> {
        let r = this.getRows()
        let c = this.getCols()
        var res = Matrix<T>(r, c)
        for (i in 0..r) {
            for (j in 0..c) {
                res[i,j] = f(this[i,j])
            }
        }
        return res
    }

    /*
     * Return the transpose of the matrix.
     */
    public func transpose(): Matrix<T> {
        let r = this.getRows()
        let c = this.getCols()
        var res = Matrix<T>(c, r)
        for (i in 0..r) {
            for (j in 0..c) {
                res[j,i] = this[i,j]
            }
        }
        return res
    }

    /*
     * Convert an n-by-1 matrix to a vector.
     */
    public func toVector(): Vector<T> {
        if (this.getCols() != 1) {
            throw IllegalArgumentException("toVector: matrix dims incompatible.")
        }

        var row = this.getRows()
        var res = Vector<T>(row)
        for (i in 0..row) {
            res[i] = this[i,0]
        }
        return res
    }

    ~init() {
        if (!ptr.isNull()) {
            unsafe { free(ptr) }
        }
    }
}

public func empty<T>(rows: Int64, cols: Int64): Matrix<T> where T <: CNumber<T> {
    return Matrix<T>(rows, cols)
}

public func zeros<T>(rows: Int64, cols: Int64): Matrix<T> where T <: CNumber<T> {
    return Matrix<T>(rows, cols, T.fromInt(0))
}

public func ones<T>(rows: Int64, cols: Int64): Matrix<T> where T <: CNumber<T> {
    return Matrix<T>(rows, cols, T.fromInt(1))
}

public func matrix<T>(rows: Int64, cols: Int64, x: T): Matrix<T> where T <: CNumber<T> {
    return Matrix<T>(rows, cols, x)
}

public func matrix<T>(rows: Int64, cols: Int64, f: (Int64, Int64) -> T): Matrix<T> where T <: CNumber<T> {
    return Matrix<T>(rows, cols, f)
}

public func matrix<T>(a: Array<Array<T>>): Matrix<T> where T <: CNumber<T> {
    return Matrix<T>(a)
}

public func matrix<T>(ptr: CPointer<Unit>, rows: Int64, cols: Int64): Matrix<T> where T <: CNumber<T> {
    return Matrix<T>(ptr, UIntNative(rows), UIntNative(cols))
}

extend<T> Matrix<T> <: ToString where T <: ToString {
    public func toString(): String {
        let builder = StringBuilder()
        builder.append("[")
        let rows = this.getRows()
        let cols = this.getCols()
        for (i in 0..rows) {
            if (i != 0) {
                builder.append(",\n ")
            }
            builder.append(this[i].toString())
        }
        builder.append("]")
        return builder.toString()
    }
}

extend<T> Matrix<T> where T <: Real<T> {
    /*
     * Convert the matrix to elements of type Float64.
     */
    public func toFloat64(): Matrix<Float64> {
        var res = Matrix<Float64>(this.getRows(), this.getCols())
        let row = this.getRows()
        let col = this.getCols()
        for (i in 0..row) {
            for (j in 0..col) {
                res[i,j] = this[i,j].toFloat64()
            }
        }
        return res
    }
}

public class Tensor<T> <: AbstractTensor<T> where T <: CNumber<T> {
    /*
     * Address storing data.
     */
    public var ptr: CPointer<Unit> = CPointer<Unit>()

    init(dims: Array<Int64>) {
        for (i in this.dims) {
            if (i <= 0) {
                throw IllegalArgumentException("init: each dimension must be positive.")
            }
        }
        this.dims = dims
        let size = this.size()
        this.ptr = unsafe { malloc(UIntNative(T.getSize() * size)) }
    }

    init(ptr: CPointer<Unit>, dims: Array<Int64>) {
        this.ptr = ptr
        this.dims = dims
    }

    init(dims: Array<Int64>, x: T) {
        for (i in this.dims) {
            if (i <= 0) {
                throw IllegalArgumentException("init: each dimension must be positive.")
            }
        }
        this.dims = dims
        let size = this.size()
        this.ptr = unsafe { malloc(UIntNative(T.getSize() * size)) }
        for (i in 0..size) {
            this.set(i, x)
        }
    }

    init(dims: Array<Int64>, f: (Array<Int64>) -> T) {
        for (i in this.dims) {
            if (i <= 0) {
                throw IllegalArgumentException("init: each dimension must be positive.")
            }
        }
        this.dims = dims
        let size = this.size()
        this.ptr = unsafe { malloc(UIntNative(T.getSize() * size)) }
        for (i in 0..size) {
            this.set(i, f(this.toIndex(i)))
        }
    }

    public func get(index: Int64): T {
        return unsafe { T.read(this.ptr, index) }
    }

    public func set(index: Int64, value: T): Unit {
        unsafe { T.write(this.ptr, value, index) }
    }

    public func copy(): Tensor<T> {
        var res = Tensor<T>(this.dims)
        let size = this.size()
        for (i in 0..size) {
            res.set(i, this.get(i))
        }
        return res
    }

    public func toIndex(i: Int64): Array<Int64> {
        let d = this.dims.size
        var index = Array<Int64>(d, {_ => 0})
        var cur_i: Int64 = i
        for (j in 0..d) {
            index[j] = cur_i % this.dims[j]
            cur_i /= this.dims[j]
        }
        return index
    }
  
    public func fromIndex(index: Array<Int64>): Int64 {
        if (this.dims.size != index.size) {
            throw IllegalArgumentException()
        }
        let d = index.size
        var i: Int64 = index[d-1]
        for (j in d-2..=0:-1) {
            i = i * this.dims[j] + index[j]
        }
        return i
    }

    /*
     * Access element of tensor by array of indices.
     */
    public operator func [](index: Array<Int64>): T {
        return this.get(this.fromIndex(index))
    }
  
    /*
     * Write to element of tensor by array of indices.
     */
    public operator func [](index: Array<Int64>, value!: T): Unit {
        this.set(this.fromIndex(index), value)
    }

    /*
     * Obtain the i'th subtensor.
     */
    public operator func [](index: Int64): Tensor<T> {
        if (this.dims.size == 1) {
            throw IllegalArgumentException("Calling subtensor on a tensor of dimension 1.")
        }
        var res = Tensor<T>(this.dims[1..])
        let size = res.size()
        for (i in 0..size) {
            res.set(i, this.get(i * this.dims[0] + index))
        }
        return res
    }

    /*
     * Convert an m-by-n tensor to a matrix
     */
    public func asMatrix(): Matrix<T> {
        if (this.dims.size != 2) {
            throw IllegalArgumentException("asMatrix: dimension is not 2.")
        }

        let r = this.dims[0]
        let c = this.dims[1]
        var res = Matrix<T>(r, c)
        for (i in 0..r) {
            for (j in 0..c) {
                res[i,j] = this[[i, j]]
            }
        }
        return res
    }

    public func apply(f: (T) -> T): Tensor<T> {
        var res = Tensor<T>(this.dims)
        var size = this.size()
        for (i in 0..size) {
            res.set(i, f(this.get(i)))
        }
        return res
    }

    ~init() {
        if (!ptr.isNull()) {
            unsafe { free(ptr) }
        }
    }
}

public func empty<T>(dims: Array<Int64>): Tensor<T> where T <: CNumber<T> {
    return Tensor<T>(dims)
}

public func zeros<T>(dims: Array<Int64>): Tensor<T> where T <: CNumber<T> {
    return Tensor<T>(dims, T.fromInt(0))
}

public func ones<T>(dims: Array<Int64>): Tensor<T> where T <: CNumber<T> {
    return Tensor<T>(dims, T.fromInt(1))
}

public func tensor<T>(dims: Array<Int64>, x: T): Tensor<T> where T <: CNumber<T> {
    return Tensor<T>(dims, x)
}

public func tensor<T>(dims: Array<Int64>, f: (Array<Int64>) -> T): Tensor<T> where T <: CNumber<T> {
    return Tensor<T>(dims, f)
}

public func tensor<T>(ptr: CPointer<Unit>, dims: Array<Int64>): Tensor<T> where T <: CNumber<T> {
    return Tensor<T>(ptr, dims)
}

extend<T> Tensor<T> where T <: Real<T> {
    /*
     * Convert the tensor to elements of type Float64.
     */
    public func toFloat64(): Tensor<Float64> {
        var res = Tensor<Float64>(this.dims)
        let size = this.size()
        for (i in 0..size) {
            res.set(i, this.get(i).toFloat64())
        }
        return res
    }
}

public interface TensorOp<T> where T <: CNumber<T> {
    /* Matrix-matrix multiplication */
    static func matprod(a: Matrix<T>, b: Matrix<T>): Matrix<T>

    /* Matrix-vector multiplication */
    static func matvecprod(a: Matrix<T>, b: Vector<T>): Vector<T>

    /* Vector-matrix multplication */
    static func vecmatprod(a: Vector<T>, b: Matrix<T>): Vector<T>

    /* Scalar-vector multiplication */
    operator func *(a: Vector<T>): Vector<T>

    /* Scalar-matrix multiplication */
    operator func *(a: Matrix<T>): Matrix<T>

    /* Scalar-tensor multiplication */
    operator func *(a: Tensor<T>): Tensor<T>
}

extend<T> Vector<T> where T <: TensorOp<T> {
    public operator func *(right: Matrix<T>): Vector<T> {
        return T.vecmatprod(this, right)
    }
}

extend<T> Matrix<T> where T <: TensorOp<T> {
    public operator func *(right: Matrix<T>): Matrix<T> {
        return T.matprod(this, right)
    }

    public operator func *(right: Vector<T>): Vector<T> {
        return T.matvecprod(this, right)
    }
}

extend Float32 <: TensorOp<Float32> {
    public static func matprod(a: Matrix<Float32>, b: Matrix<Float32>): Matrix<Float32> {
        return cj_sgemm(a, b)
    }

    public static func matvecprod(a: Matrix<Float32>, b: Vector<Float32>): Vector<Float32> {
        return cj_sgemv(a, b)
    }

    public static func vecmatprod(a: Vector<Float32>, b: Matrix<Float32>): Vector<Float32> {
        return cj_sgemv_trans(a, b)
    }

    public operator func *(a: Vector<Float32>): Vector<Float32> {
        return a.apply({x => this * x})
    }

    public operator func *(a: Matrix<Float32>): Matrix<Float32> {
        return a.apply({x => this * x})
    }

    public operator func *(a: Tensor<Float32>): Tensor<Float32> {
        return a.apply({x => this * x})
    }
}

extend Float64 <: TensorOp<Float64> {
    public static func matprod(a: Matrix<Float64>, b: Matrix<Float64>): Matrix<Float64> {
        return cj_dgemm(a, b)
    }

    public static func matvecprod(a: Matrix<Float64>, b: Vector<Float64>): Vector<Float64> {
        return cj_dgemv(a, b)
    }

    public static func vecmatprod(a: Vector<Float64>, b: Matrix<Float64>): Vector<Float64> {
        return cj_dgemv_trans(a, b)
    }

    public operator func *(a: Vector<Float64>): Vector<Float64> {
        return a.apply({x => this * x})
    }

    public operator func *(a: Matrix<Float64>): Matrix<Float64> {
        return a.apply({x => this * x})
    }

    public operator func *(a: Tensor<Float64>): Tensor<Float64> {
        return a.apply({x => this * x})
    }
}

extend Complex64 <: TensorOp<Complex64> {
    public static func matprod(a: Matrix<Complex64>, b: Matrix<Complex64>): Matrix<Complex64> {
        return cj_zgemm(a, b)
    }

    public static func matvecprod(a: Matrix<Complex64>, b: Vector<Complex64>): Vector<Complex64> {
        return cj_zgemv(a, b)
    }

    public static func vecmatprod(a: Vector<Complex64>, b: Matrix<Complex64>): Vector<Complex64> {
        return cj_zgemv_trans(a, b)
    }

    public operator func *(a: Vector<Complex64>): Vector<Complex64> {
        return a.apply({x => this * x})
    }

    public operator func *(a: Matrix<Complex64>): Matrix<Complex64> {
        return a.apply({x => this * x})
    }

    public operator func *(a: Tensor<Complex64>): Tensor<Complex64> {
        return a.apply({x => this * x})
    }
}

extend Complex32 <: TensorOp<Complex32> {
    public static func matprod(a: Matrix<Complex32>, b: Matrix<Complex32>): Matrix<Complex32> {
        return cj_cgemm(a, b)
    }

    public static func matvecprod(a: Matrix<Complex32>, b: Vector<Complex32>): Vector<Complex32> {
        return cj_cgemv(a, b)
    }

    public static func vecmatprod(a: Vector<Complex32>, b: Matrix<Complex32>): Vector<Complex32> {
        return cj_cgemv_trans(a, b)
    }

    public operator func *(a: Vector<Complex32>): Vector<Complex32> {
        return a.apply({x => this * x})
    }

    public operator func *(a: Matrix<Complex32>): Matrix<Complex32> {
        return a.apply({x => this * x})
    }

    public operator func *(a: Tensor<Complex32>): Tensor<Complex32> {
        return a.apply({x => this * x})
    }
}

public func approxEqual<T>(a: Vector<T>, b: Vector<T>, atol!:Float64 = 1e-7): Bool where T <: Float<T> & ToString {
    if (a.size() != b.size()) {
        return false
    }
    var size = a.size()
    for (i in 0..size) {
        if (!approxEqual<T>(a[i], b[i], atol:atol)) {
            return false
        }
    }
    return true
}

public func assertApproxEqual<T>(a: Vector<T>, b: Vector<T>, atol!:Float64 = 1e-7): Unit where T <: Float<T> & ToString {
    if (a.size() != b.size()) {
        throw AssertionException("assertApproxEqual: vector dimension is not same.")
    }
    var size = a.size()
    for (i in 0..size) {
        assertApproxEqual<T>(a[i], b[i], atol:atol)
    }
}

public func approxEqualC<T>(a: Vector<T>, b: Vector<T>, atol!:Float64 = 1e-7): Bool where T <: Complex<T> & ToString {
    if (a.size() != b.size()) {
        return false
    }
    var size = a.size()
    for (i in 0..size) {
        if (!approxEqualC<T>(a[i], b[i], atol:atol)) {
            return false
        }
    }
    return true
}

public func assertApproxEqualC<T>(a: Vector<T>, b: Vector<T>, atol!:Float64 = 1e-7): Unit where T <: Complex<T> & ToString {
    if (a.size() != b.size()) {
        throw AssertionException("assertApproxEqual: vector dimension is not same.")
    }
    var size = a.size()
    for (i in 0..size) {
        assertApproxEqualC<T>(a[i], b[i], atol:atol)
    }
}

public func approxEqual<T>(a: Matrix<T>, b: Matrix<T>, atol!:Float64 = 1e-7): Bool where T <: Float<T> & ToString {
    if (a.getRows() != b.getRows() || a.getCols() != b.getCols()) {
        return false
    }
    var row = a.getRows()
    var col = a.getCols()
    for (i in 0..row) {
        for (j in 0..col) {
            if (!approxEqual<T>(a[i,j], b[i,j], atol:atol)) {
                return false
            }
        }
    }
    return true
}

public func assertEqual<T>(a: Matrix<T>, b: Matrix<T>): Unit where T <: CInteger<T> & ToString {
    if (a.getRows() != b.getRows() || a.getCols() != b.getCols()) {
        throw AssertionException("assertApproxEqual: matrix dimension is not same.")
    }
    var row = a.getRows()
    var col = a.getCols()
    for (i in 0..row) {
        for (j in 0..col) {
            assertEqual<T>(a[i,j], b[i,j])
        }
    }
}

public func assertApproxEqual<T>(a: Matrix<T>, b: Matrix<T>, atol!:Float64 = 1e-7): Unit where T <: Float<T> & ToString {
    if (a.getRows() != b.getRows() || a.getCols() != b.getCols()) {
        throw AssertionException("assertApproxEqual: matrix dimension is not same.")
    }
    var row = a.getRows()
    var col = a.getCols()
    for (i in 0..row) {
        for (j in 0..col) {
            assertApproxEqual<T>(a[i,j], b[i,j], atol:atol)
        }
    }
}

public func approxEqualC<T>(a: Matrix<T>, b: Matrix<T>, atol!:Float64 = 1e-7): Bool where T <: Complex<T> & ToString {
    if (a.getRows() != b.getRows() || a.getCols() != b.getCols()) {
        return false
    }
    var row = a.getRows()
    var col = a.getCols()
    for (i in 0..row) {
        for (j in 0..col) {
            if (!approxEqualC<T>(a[i,j], b[i,j], atol:atol)) {
                return false
            }
        }
    }
    return true
}

public func assertApproxEqualC<T>(a: Matrix<T>, b: Matrix<T>, atol!:Float64 = 1e-7): Unit where T <: Complex<T> & ToString {
    if (a.getRows() != b.getRows() || a.getCols() != b.getCols()) {
        throw AssertionException("assertApproxEqual: matrix dimension is not same.")
    }
    var row = a.getRows()
    var col = a.getCols()
    for (i in 0..row) {
        for (j in 0..col) {
            assertApproxEqualC<T>(a[i,j], b[i,j], atol:atol)
        }
    }
}

public func arange<T>(start: T, stop: T, step: T): Vector<T> where T <: Real<T> & Comparable<T> {
    var res = ArrayList<T>()
    var x: T = start
    while (x < stop) {
        res.add(x)
        x = x + step
    }
    return Vector<T>(res.toArray())
}

public func linspace<T>(start: T, stop: T, num!:Int64 = 100): Vector<T> where T <: Float<T> {
    var res = Vector<T>(num)
    var diff = (stop - start) / T.fromInt(num - 1)
    for (i in 0..num) {
        res[i] = start + diff * T.fromInt(i)
    }
    return res
}

public func logspace<T>(start: T, stop: T, num!:Int64 = 50): Vector<T> where T <: Float<T> {
    var res = linspace<T>(start, stop, num: num)
    return res.apply({x => power(T.fromInt(10), x)})
}

public func geomspace<T>(start: T, stop: T, num!:Int64 = 50): Vector<T> where T <: Float<T> {
    var res = linspace<T>(log10(start), log10(stop), num: num)
    return res.apply({x => power(T.fromInt(10), x)})
}

public func meshgrid<T>(x: Vector<T>, y: Vector<T>): (Matrix<T>, Matrix<T>) where T <: CNumber<T> {
    let row = y.size()
    let col = x.size()
    let X = Matrix<T>(row, col)
    let Y = Matrix<T>(row, col)
    for (i in 0..row) {
        for (j in 0..col) {
            X[i,j] = x[j]
            Y[i,j] = y[i]
        }
    }
    return (X, Y)
}

public func meshgrid<T>(x: Vector<T>, y: Vector<T>, f: (T, T) -> T):
    (Matrix<T>, Matrix<T>, Matrix<T>) where T <: CNumber<T> {
    let row = y.size()
    let col = x.size()
    let X = Matrix<T>(row, col)
    let Y = Matrix<T>(row, col)
    let Z = Matrix<T>(row, col)
    for (i in 0..row) {
        for (j in 0..col) {
            X[i,j] = x[j]
            Y[i,j] = y[i]
            Z[i,j] = f(x[j], y[i])
        }
    }
    return (X, Y, Z)
}

public func apply2<T>(x: Matrix<T>, y: Matrix<T>, f: (T, T) -> T): Matrix<T> where T <: CNumber<T> {
    if (x.shape() != y.shape()) {
        throw IllegalArgumentException("apply2: dimension mismatch")
    }

    let row = x.getRows()
    let col = y.getCols()
    let res = Matrix<T>(row, col)
    for (i in 0..row) {
        for (j in 0..col) {
            res[i,j] = f(x[i,j], y[i,j])
        }
    }
    return res
}

public func apply2<T>(x: Vector<T>, y: Vector<T>, f: (T, T) -> T): Vector<T> where T <: CNumber<T> {
    if (x.size() != y.size()) {
        throw IllegalArgumentException("apply2: dimension mismatch")
    }

    let size = x.size()
    let res = Vector<T>(size)
    for (i in 0..size) {
        res[i] = f(x[i], y[i])
    }
    return res
}

/*
 * Compute estimated gradient on a vector using centered differencing.
 */
public func gradient<T>(v: Vector<T>, spacing!: T): Vector<T> where T <: Float<T> {
    let size = v.size()
    let res = empty<T>(size)
    res[0] = (v[1] - v[0]) / spacing
    res[size-1] = (v[size-1] - v[size-2]) / spacing
    for (i in 1..size-1) {
        res[i] = (v[i+1] - v[i-1]) / spacing / T.fromInt(2)
    }
    return res
}

/*
 * Compute estimated gradient on a matrix using centered differencing.
 */
public func gradient<T>(m: Matrix<T>, spacing_x!: T, spacing_y!: T): (Matrix<T>, Matrix<T>) where T <: Float<T> {
    let r = m.getRows()
    let c = m.getCols()
    let dx = empty<T>(r, c)
    for (i in 0..r) {
        // Compute for each row
        dx[i,0] = (m[i,1] - m[i,0]) / spacing_x
        dx[i,c-1] = (m[i,c-1] - m[i,c-2]) / spacing_x
        for (j in 1..c-1) {
            dx[i,j] = (m[i,j+1] - m[i,j-1]) / spacing_x / T.fromInt(2)
        }
    }
    let dy = empty<T>(r, c)
    for (j in 0..c) {
        // Compute for each column
        dy[0,j] = (m[1,j] - m[0,j]) / spacing_y
        dy[r-1,j] = (m[r-1,j] - m[r-2,j]) / spacing_y
        for (i in 1..r-1) {
            dy[i,j] = (m[i+1,j] - m[i-1,j]) / spacing_y / T.fromInt(2)
        }
    }
    return (dx, dy)
}

public func sorted<T>(v: Vector<T>, stable!:Bool = false): Vector<T> where T <: Real<T> & Comparable<T> {
    let a: Array<T> = v.toArray()
    sort(a)
    return Vector<T>(a)
}

public func sortedDescending<T>(v: Vector<T>, stable!:Bool = false): Vector<T> where T <: Real<T> & Comparable<T> {
    let a: Array<T> = v.toArray()
    sort(a, descending: true)
    return Vector<T>(a)
}

public func concat<T>(left: Vector<T>, right: Vector<T>): Vector<T> where T <: CNumber<T> {
    let lsz = left.size()
    let rsz = right.size()
    let sz = lsz + rsz
    let x = Vector<T>(sz)
    for (i in 0..sz) {
        if (i < lsz) {
            x[i] = left[i]
        }
        else {
            x[i] = right[i-lsz]
        }
    }
    return x
}

public func concat<T>(a: Array<Vector<T>>): Vector<T> where T <: CNumber<T> {
    var size: Int64 = 0
    for (i in 0..a.size) {
        size += a[i].size()
    }
    var res = Vector<T>(size)
    var pos: Int64 = 0
    for (i in 0..a.size) {
        for (j in 0..a[i].size()) {
            res[pos] = a[i][j]
            pos++
        }
    }
    return res
}

public func eye<T>(size: Int64): Matrix<T> where T <: CNumber<T> {
    let res = zeros<T>(size, size)
    for (i in 0..size) {
        res[i,i] = T.fromInt(1) 
    }
    return res
}

public func eye<T>(row: Int64, col: Int64): Matrix<T> where T <: CNumber<T> {
    let res = zeros<T>(row, col)
    for (i in 0..min(row, col)) {
        res[i,i] = T.fromInt(1)
    }
    return res
}

public func diag<T>(a: Vector<T>): Matrix<T> where T <: CNumber<T> {
    let n = a.size()
    let res = zeros<T>(n, n)
    for (i in 0..n) {
        res[i,i] = a[i]
    }
    return res
}

public func diag<T>(a: Vector<T>, m: Int64, n: Int64): Matrix<T> where T <: CNumber<T> {
    let res = zeros<T>(m, n)
    for (i in 0..m) {
        for (j in 0..n) {
            res[i,i] = a[i]
        }   
    }
    return res
}

public func tri<T>(rows: Int64, cols: Int64, k!:Int64 = 0): Matrix<T> where T <: CNumber<T> {
    let res = zeros<T>(rows, cols)
    for (i in 0..rows) {
        for (j in 0..cols) {
            if (j - i <= k) {
                res[i,j] = T.fromInt(1)
            }
        }
    }
    return res
}

public func tril<T>(a: Matrix<T>, k!: Int64 = 0): Matrix<T> where T <: CNumber<T> {
    let rows = a.getRows()
    let cols = a.getCols()
    let res = zeros<T>(rows, cols)
    for (i in 0..rows) {
        for (j in 0..cols) {
            if (j - i <= k) {
                res[i,j] = a[i,j]
            }
        }
    }
    return res
}

public func triu<T>(a: Matrix<T>, k!: Int64 = 0): Matrix<T> where T <: CNumber<T> {
    let rows = a.getRows()
    let cols = a.getCols()
    let res = zeros<T>(rows, cols)
    for (i in 0..rows) {
        for (j in 0..cols) {
            if (j - i >= k) {
                res[i,j] = a[i,j]
            }
        }
    }
    return res
}

public func trace<T>(a: Matrix<T>): T where T <: CNumber<T> {
    if (a.getRows() != a.getCols()) {
        throw IllegalArgumentException("trace: dimension mismatch")
    }
    
    let n = a.dims[0]
    var res = T.fromInt(0)
    for (i in 0..n) {
        res = res + a[i,i]
    }
    return res
}

/*
 * Flatten a matrix by row if order = true, by column if order = false.
 */
public func flatten<T>(a: Matrix<T>, order!:Bool = true): Vector<T> where T <: CNumber<T> {
    let n = a.size()
    let res = zeros<T>(n)
    let row = a.getRows()
    let col = a.getCols()

    if (order) {
        for (i in 0..row) {
            for (j in 0..col) {
                res[i * col + j] = a[i,j]
            }
        }
    } else {
        for (j in 0..col) {
            for (i in 0..row) {
                res[j * row + i] = a[i,j]
            }
        }
    }

    return res
}

/*
 * Reshape matrix to vector.
 */
public func reshape<T>(a: Matrix<T>): Vector<T> where T <: CNumber<T> {
    return flatten(a)
}

/*
 * Reshape vector to matrix.
 */
public func reshape<T>(a: Vector<T>, row: Int64, col: Int64): Matrix<T> where T <: CNumber<T> {
    if (a.size() != row * col) {
        throw IllegalArgumentException("reshape: number of elements in a is not row * col.")
    }

    let res = empty<T>(row, col)
    for (i in 0..row) {
        for (j in 0..col) {
            res[i,j] = a[i*col+j]
        }
    }
    return res
}

/*
 * Compute the Kronecker product
 */
public func kron<T>(a: Matrix<T>, b: Matrix<T>): Matrix<T> where T <: CNumber<T> {
    let row_a = a.dims[0]
    let col_a = a.dims[1]
    let row_b = b.dims[0]
    let col_b = b.dims[1]

    let row_res = row_a * row_b
    let col_res = col_a * col_b

    let res = zeros<T>(row_res, col_res)
    for (i in 0..row_res) {
        for (j in 0..col_res) {
            res[i,j] = a[i / row_b, j / col_b] * b[i % row_b, j % col_b]
        }
    }

    return res
}

public func sum<T>(a: Vector<T>): T where T <: CNumber<T> {
    let size = a.size()
    var res = T.fromInt(0)
    for (i in 0..size) {
        res += a[i]
    }
    return res
}

public func sumRow<T>(a: Matrix<T>): Vector<T> where T <: CNumber<T> {
    let row = a.getRows()
    let col = a.getCols()
    var res = zeros<T>(row)
    for (i in 0..row) {
        for (j in 0..col) {
            res[i] += a[i,j]
        }
    }
    return res
}

public func sumCol<T>(a: Matrix<T>): Vector<T> where T <: CNumber<T> {
    let row = a.getRows()
    let col = a.getCols()
    var res = zeros<T>(col)
    for (i in 0..row) {
        for (j in 0..col) {
            res[j] += a[i,j]
        }
    }
    return res
}

/*
 *  Compute the minimum.
 */
public func min<T>(a: Vector<T>): T where T <: Real<T> & Comparable<T> {
    var res = a[0]
    let n = a.size()
    for (i in 1..n) {
        if (res > a[i]) {
            res = a[i]
        }
    }
    return res 
}

/*
 *  Compute the maximum.
 */
public func max<T>(a: Vector<T>): T where T <: Real<T> & Comparable<T> {
    var res = a[0]
    let n = a.size()
    for (i in 1..n) {
        if (res < a[i]) {
            res = a[i]
        }
    }
    return res 
}

public func conjugate<T>(a: Vector<T>): Vector<T> where T <: Complex<T> {
    return a.apply({x => x.conjugate()})
}

public func conjugate<T>(a: Matrix<T>): Matrix<T> where T <: Complex<T> {
    return a.apply({x => x.conjugate()})
}

public func htranspose<T>(a: Matrix<T>): Matrix<T> where T <: Complex<T> {
    let r = a.getRows()
    let c = a.getCols()
    return matrix<T>(c, r, {i:Int64, j:Int64 => a[j,i].conjugate()})
}

public func hstack<T>(a: Matrix<T>, b: Matrix<T>): Matrix<T> where T <: CNumber<T> {
    let row_a = a.getRows()
    let col_a = a.getCols()
    let row_b = b.getRows()
    let col_b = b.getCols()
    let row_res = row_a
    let col_res = col_a + col_b
    let res = zeros<T>(row_res, col_res)
    for (i in 0..row_res) {
        for (j in 0..col_res) {
            if (j < col_a) {
                res[i,j] = a[i,j]
            } else {
                res[i,j] = b[i,j-col_a]
            }
        }
    }
    return res
}

public func vstack<T>(a: Matrix<T>, b: Matrix<T>): Matrix<T> where T <: CNumber<T> {
    let row_a = a.getRows()
    let col_a = a.getCols()
    let row_b = b.getRows()
    let col_b = b.getCols()
    let row_res = row_a + row_b
    let col_res = col_a
    let res = zeros<T>(row_res, col_res)
    for (i in 0..row_res) {
        for (j in 0..col_res) {
            if (i < row_a) {
                res[i,j] = a[i,j]
            } else {
                res[i,j] = b[i-row_a,j]
            }
        }
    }
    return res
}

@Test
public class TestTensor {
    @TestCase
    func testVectorBasic(): Unit {
        var a1: Vector<Float32> = empty(10)
        @Assert(a1.dtype() == "Float32")
        @Assert(a1.shape() == [10])
        @Assert(a1.size() == 10)

        var a2: Vector<Float32> = vector(10, {i: Int64 => Float32(i)})
        @Assert(approxEqual(a2, vector<Float32>(
            [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]
        )))
    }

    @TestCase
    func testMatrixBasic(): Unit {
        var a1: Matrix<Float64> = empty(5, 5)
        @Assert(a1.dtype() == "Float64")
        @Assert(a1.shape() == [5, 5])
        @Assert(a1.size() == 25)

        var a2: Matrix<Float64> = matrix(
            5, 5, {i: Int64, j: Int64 => Float64(i * 10 + j)}
        )
        @Assert(approxEqual(a2, matrix<Float64>(
            [[ 0.0,  1.0,  2.0,  3.0,  4.0],
             [10.0, 11.0, 12.0, 13.0, 14.0],
             [20.0, 21.0, 22.0, 23.0, 24.0],
             [30.0, 31.0, 32.0, 33.0, 34.0],
             [40.0, 41.0, 42.0, 43.0, 44.0]]
        )))
    }

    @TestCase
    func testTensorBasic(): Unit {
        var a1: Tensor<Float32> = empty([2, 2, 2])
        @Assert(a1.dtype() == "Float32")
        @Assert(a1.shape() == [2, 2, 2])
        @Assert(a1.size() == 8)

        var a2: Tensor<Float32> = tensor(
            [2, 2, 2], {id: Array<Int64> => Float32(id[0] * 100 + id[1] * 10 + id[2])}
        )
        @Assert(approxEqual(a2[0].asMatrix(), matrix<Float32>(
            [[ 0.0,  1.0],
            [10.0, 11.0]]
        )))
        @Assert(approxEqual(a2[1].asMatrix(), matrix<Float32>(
            [[100.0, 101.0],
            [110.0, 111.0]]
        )))
    }

    @TestCase
    func testRangeIndex(): Unit {
        var a1 = vector<Int64>([1, 3, 5, 7, 9])
        @Assert(a1[1..3] == vector<Int64>([3, 5]))
        @Assert(a1[..3] == vector<Int64>([1, 3, 5]))
        @Assert(a1[1..] == vector<Int64>([3, 5, 7, 9]))
    }

    @TestCase
    func testApply(): Unit {
        var ls1 = vector<Float64>(
            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
        )
        var ls2 = ls1.apply({t => sin(t)})
        @Assert(approxEqual(ls2, vector<Float64>(
            [0.000000, 0.099833, 0.198669, 0.295520, 0.389418, 0.479426,
            0.564642, 0.644218, 0.717356, 0.783327, 0.841471]
        ), atol:1e-6))

        var ls3 = ls1.apply({a: Float64 => 2.0 * a})
        @Assert(approxEqual(ls3, vector<Float64>(
            [0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0]
        )))
    }

    @TestCase
    func testLinspace(): Unit {
        var ls1 = linspace(Float64(0.0), 1.0, num:11)
        @Assert(approxEqual(ls1, vector<Float64>(
            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
        )))
    }

    @TestCase
    func testLogspace(): Unit {
        var a1 = logspace<Float64>(-1.0, 2.0, num:4)
        @Assert(approxEqual(a1, vector([0.1, 1.0, 10.0, 100.0])))
    }

    @TestCase
    func testGeomspace(): Unit {
        var a1 = geomspace<Float64>(0.1, 100.0, num:4)
        @Assert(approxEqual(a1, vector([0.1, 1.0, 10.0, 100.0])))
    }

    @TestCase
    func testMeshgrid(): Unit {
        var x = linspace(0.0, 1.0, num:3)
        var y = linspace(0.0, 1.0, num:2)
        var (X, Y) = meshgrid(x, y)
        @Assert(approxEqual(X, matrix<Float64>([
            [0.0, 0.5, 1.0],
            [0.0, 0.5, 1.0]
        ])))
        @Assert(approxEqual(Y, matrix<Float64>([
            [0.0, 0.0, 0.0],
            [1.0, 1.0, 1.0]
        ])))
    }

    @TestCase
    func testMeshgrid2(): Unit {
        var x: Vector<Float64> = linspace(0.0, 1.0, num:3)
        var y: Vector<Float64> = linspace(0.0, 1.0, num:2)
        var (X, Y, Z) = meshgrid(x, y, {s:Float64, t:Float64 => s + t})
        @Assert(approxEqual(X, matrix<Float64>([
            [0.0, 0.5, 1.0],
            [0.0, 0.5, 1.0]
        ])))
        @Assert(approxEqual(Y, matrix<Float64>([
            [0.0, 0.0, 0.0],
            [1.0, 1.0, 1.0]
        ])))
        @Assert(approxEqual(Z, matrix<Float64>([
            [0.0, 0.5, 1.0],
            [1.0, 1.5, 2.0]
        ])))
    }

    @TestCase
    func testGradient(): Unit {
        let x = vector([1.0, 4.0, 9.0])
        let dx = gradient(x, spacing: 1.0)
        @Assert(approxEqual(dx, vector([3.0, 4.0, 5.0])))
    }

    @TestCase
    func testGradient2(): Unit {
        let m = matrix([[1.0, 4.0, 9.0],
                        [4.0, 9.0, 16.0],
                        [9.0, 16.0, 25.0]])
        let (dx, dy) = gradient(m, spacing_x: 1.0, spacing_y: 1.0)
        @Assert(approxEqual(dx, matrix([
            [3.0, 4.0, 5.0],
            [5.0, 6.0, 7.0],
            [7.0, 8.0, 9.0]
        ])))
        @Assert(approxEqual(dy, matrix([
            [3.0, 5.0, 7.0],
            [4.0, 6.0, 8.0],
            [5.0, 7.0, 9.0]
        ])))
    }

    @TestCase
    func testTri(): Unit {
        let X = tri<Int64>(3, 5, k: 2)
        @Assert(X == matrix<Int64>([
            [1, 1, 1, 0, 0],
            [1, 1, 1, 1, 0],
            [1, 1, 1, 1, 1]
        ]))

        let Y = tri<Float64>(3, 5, k: -1)
        @Assert(Y == matrix<Float64>([
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0, 0.0, 0.0],
            [1.0, 1.0, 0.0, 0.0, 0.0]
        ]))
    }

    @TestCase
    func testTril(): Unit {
        let X = tril(matrix<Int64>([
            [1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]
        ]), k: -1)
        @Assert(X == matrix<Int64>([
            [0, 0, 0], [4, 0, 0], [7, 8, 0], [10, 11, 12]
        ]))
    }

    @TestCase
    func testTriu(): Unit {
        let X = triu(matrix<Int64>([
            [1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]
        ]), k: -1)
        @Assert(X == matrix<Int64>([
            [1, 2, 3], [4, 5, 6], [0, 8, 9], [0, 0, 12]
        ]))
    }

    @TestCase
    func testTrace(): Unit {
        let X = matrix<Int64>([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
        @Assert(trace(X) == 15)

        let Y = matrix<Float64>([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
        @Assert(approxEqual(trace(Y), 15.0))
    }

    @TestCase
    func testReshape(): Unit {
        let X = vector<Float64>([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
        let expected_res = matrix<Float64>([
            [1.0, 2.0, 3.0],
            [4.0, 5.0, 6.0]
        ])
        @Assert(approxEqual(reshape(X, 2, 3), expected_res))
    }

    @TestCase
    func testReshape2(): Unit {
        let X = matrix<Float64>([
            [1.0, 2.0, 3.0],
            [4.0, 5.0, 6.0]
        ])
        let expected_res = vector<Float64>([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
        @Assert(approxEqual(reshape(X), expected_res))
    }

    // @TestCase
    // func testDotProduct(): Unit {
    //     var a1 = vector<Int64>([1, 2, 3, 4])
    //     var a2 = vector<Int64>([2, 3, -1, 0])
    //     @Assert(a1.dot(a2) == 5)

    //     var a3 = vector<Float64>([1.0, 2.0, 3.0, 4.0])
    //     var a4 = vector<Float64>([2.0, 3.0, -1.0, 0.0])
    //     @Assert(approxEqual(a3.dot(a4), 5.0))
    // }

    @TestCase
    func testVectorNorm(): Unit {
        var a1 = vector<Float64>([3.0, 4.0])
        @Assert(approxEqual(a1.norm(), 5.0))
    }

    @TestCase
    func testSort(): Unit {
        var a1 = vector<Int64>([1, 5, 3, 9, 7])
        var a2 = sorted(a1)
        @Assert(a2 == vector<Int64>([1, 3, 5, 7, 9]))
        var a3 = sortedDescending(a1)
        @Assert(a3 == vector<Int64>([9, 7, 5, 3, 1]))

        var b1 = vector<Float64>([1.0, 5.0, 3.0, 9.0, 7.0])
        var b2 = sorted(b1)
        @Assert(b2 == vector<Float64>([1.0, 3.0, 5.0, 7.0, 9.0]))
        var b3 = sortedDescending(b1)
        @Assert(b3 == vector<Float64>([9.0, 7.0, 5.0, 3.0, 1.0]))
    }

    @TestCase
    func testComplexVector(): Unit {
        let a = vector<Complex64>([
            Complex64(1.0, 2.0), Complex64(2.0, 1.0)
        ])
        let b = vector<Complex64>([
            Complex64(1.0, 1.0), Complex64(2.0, 2.0)
        ])
        @Assert(approxEqualC(a + b, vector<Complex64>([
            Complex64(2.0, 3.0), Complex64(4.0, 3.0)
        ])))
        @Assert(approxEqualC(conjugate(a), vector<Complex64>([
            Complex64(1.0, -2.0), Complex64(2.0, -1.0)
        ])))
    }

    @TestCase
    func testComplexMatrix(): Unit {
        let a = matrix<Complex64>([
            [Complex64(1.0, 2.0), Complex64(2.0, 1.0)],
            [Complex64(1.0, 1.0), Complex64(2.0, 2.0)]
        ])
        let expected_conjugate = matrix<Complex64>([
            [Complex64(1.0, -2.0), Complex64(2.0, -1.0)],
            [Complex64(1.0, -1.0), Complex64(2.0, -2.0)]
        ])
        let expected_htranspose = matrix<Complex64>([
            [Complex64(1.0, -2.0), Complex64(1.0, -1.0)],
            [Complex64(2.0, -1.0), Complex64(2.0, -2.0)]
        ])
        @Assert(approxEqualC(conjugate(a), expected_conjugate))
        @Assert(approxEqualC(htranspose(a), expected_htranspose))
    }
}