mat.bend source
mat.bend on the hub · documented module
# bend-blas: the shaped layer. A Mat<rows, cols> is one row-major# Array<F32> with its shape in the type, a Vec<n> likewise, and every# routine below states the shape rule in its signature, so a wrong# operand is a checker error naming both shapes, not a wrong answer:# Mat.gemm(rows, k, cols, alpha, a, b, beta, c) needs a Mat<rows, k>, a# Mat<k, cols> and a Mat<rows, cols>; a transpose is a separate def# (gemm_t, gemm_tn, gemv_t), since it changes the shape.# The sizes are Nat arguments that double as the type indices. The raw# routines of blas.bend and lapack.bend run underneath, unchanged; the# shaped forms answer the same Result, with .try twins, and consume# their operands the same way (Mat.clone keeps a copy). This is the# scipy.linalg to blas.bend's scipy.linalg.blas, plus the shapes.import Baseimport ./blas.bend as Bimport ./lapack.bend as Ltype Mat<-rows: Nat, -cols: Nat> is Type: MatA{a: Array<F32>}type Vec<-n: Nat> is Type: VecA{a: Array<F32>}# Construction# ------------# the depth of the smallest power-of-two block holding count floatsdef depth.go(fuel: Nat, +n: U32, +d: Nat, +cap: U32) -> Nat: match fuel: case 0n: d case 1n+p: Bool.pick(Nat, U32.is_lt(cap, n), depth.go(p, n, 1n+d, (cap * 2 : U32)), d)def depth(+count: Nat) -> Nat: depth.go(40n, U32.from_nat(count), 0n, 1)def fill.go(xs: List<&2, F32>, +i: U32, a: Array<F32>) -> Array<F32>: match xs: case []: a case x <> rest: fill.go(rest, (i + 1 : U32), Array.set(F32, a, i, x))# an array of the smallest power-of-two size holding count floats,# the first ones from xs (cut or padded with 0.0), the rest 0.0def array.of(+count: Nat, xs: List<&2, F32>) -> Array<F32>: fill.go(List.take(&2, F32, xs, count), 0, Array.new(F32, depth(count), 0.0))def take.go(n: Nat, +i: U32, r: Array<F32> & F32, acc: List<&2, F32>) -> List<&2, F32>: match n: case 0n: (a, x) = r x <> acc case 1n+p: (a, x) = r take.go(p, (i - 1 : U32), Array.get(F32, a, (i - 1 : U32)), x <> acc)# the first count floats of an array, as a listdef list.of(+count: Nat, a: Array<F32>) -> List<&2, F32>: match count: case 0n: [] case 1n+p: take.go(p, (U32.from_nat(p) : U32), Array.get(F32, a, U32.from_nat(p)), [])# xs cut into rows of colsdef rows.of(rows: Nat, +cols: Nat, +xs: List<&2, F32>) -> List<&2, List<&2, F32>>: match rows: case 0n: [] case 1n+p: List.take(&2, F32, xs, cols) <> rows.of(p, cols, List.drop(&2, F32, xs, cols))def F32.show.of(x: F32) -> String: F32.show(x)def row.show(xs: List<&2, F32>) -> String: List.show(~&2, ~F32, ~F32.show.of, xs)# rows * cols floats, row-major, cut or padded with 0.0def Mat.from_list(+rows: Nat, +cols: Nat, xs: List<&2, F32>) -> Mat<rows, cols>: MatA{array.of(Nat.mul(rows, cols), xs)}def Mat.zeros(+rows: Nat, +cols: Nat) -> Mat<rows, cols>: MatA{Array.new(F32, depth(Nat.mul(rows, cols)), 0.0)}def Mat.to_list(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, F32>: MatA{a} = m list.of(Nat.mul(rows, cols), a)def Mat.to_rows(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> List<&2, List<&2, F32>>: rows.of(rows, cols, Mat.to_list(rows, cols, m))def Mat.show(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> String: List.show(~&2, ~List<&2, F32>, ~row.show, Mat.to_rows(rows, cols, m))def Mat.clone.fin(-rows: Nat, -cols: Nat, r: Array<F32> & Array<F32>) -> Mat<rows, cols> & Mat<rows, cols>: (x, y) = r (MatA{x}, MatA{y})# a Mat is affine: two copies, to be opened by matching in a def of its owndef Mat.clone(+rows: Nat, +cols: Nat, m: Mat<rows, cols>) -> Mat<rows, cols> & Mat<rows, cols>: MatA{a} = m Mat.clone.fin(rows, cols, Array.clone(F32, a))def Vec.from_list(+n: Nat, xs: List<&2, F32>) -> Vec<n>: VecA{array.of(n, xs)}def Vec.zeros(+n: Nat) -> Vec<n>: VecA{Array.new(F32, depth(n), 0.0)}def Vec.to_list(+n: Nat, v: Vec<n>) -> List<&2, F32>: VecA{a} = v list.of(n, a)def Vec.show(+n: Nat, v: Vec<n>) -> String: row.show(Vec.to_list(n, v))def Vec.clone.fin(-n: Nat, r: Array<F32> & Array<F32>) -> Vec<n> & Vec<n>: (x, y) = r (VecA{x}, VecA{y})def Vec.clone(+n: Nat, v: Vec<n>) -> Vec<n> & Vec<n>: VecA{a} = v Vec.clone.fin(n, Array.clone(F32, a))# the raw answers, shapeddef Mat.wrap(-rows: Nat, -cols: Nat, r: Result<&1, &1, U32 & String, Array<F32>>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>): IO.pure(Result<&1, &1, U32 & String, Mat<rows, cols>>, Result.map(&1, &1, U32 & String, Array<F32>, Mat<rows, cols>, x => MatA{x}, r))def Vec.wrap(-n: Nat, r: Result<&1, &1, U32 & String, Array<F32>>) -> IO(Result<&1, &1, U32 & String, Vec<n>>): IO.pure(Result<&1, &1, U32 & String, Vec<n>>, Result.map(&1, &1, U32 & String, Array<F32>, Vec<n>, x => VecA{x}, r))# BLAS# ----# C = alpha A B + beta Cdef Mat.gemm(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>): MatA{x} = a MatA{y} = b MatA{z} = c IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(0, 0, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))def Mat.gemm.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>): IO.try(Mat<rows, cols>, Mat.gemm(rows, k, cols, alpha, a, b, beta, c))# C = alpha A B^T + beta C, B stored cols x kdef Mat.gemm_t(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<cols, k>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>): MatA{x} = a MatA{y} = b MatA{z} = c IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(0, 1, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))def Mat.gemm_t.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<rows, k>, b: Mat<cols, k>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>): IO.try(Mat<rows, cols>, Mat.gemm_t(rows, k, cols, alpha, a, b, beta, c))# C = alpha A^T B + beta C, A stored k x rowsdef Mat.gemm_tn(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<k, rows>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>): MatA{x} = a MatA{y} = b MatA{z} = c IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sgemm(1, 0, U32.from_nat(rows), U32.from_nat(cols), U32.from_nat(k), alpha, x, y, beta, z), r => Mat.wrap(rows, cols, r))def Mat.gemm_tn.try(+rows: Nat, +k: Nat, +cols: Nat, alpha: F32, a: Mat<k, rows>, b: Mat<k, cols>, beta: F32, c: Mat<rows, cols>) -> IO(Mat<rows, cols>): IO.try(Mat<rows, cols>, Mat.gemm_tn(rows, k, cols, alpha, a, b, beta, c))# y = alpha A x + beta ydef Mat.gemv(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<cols>, beta: F32, y: Vec<rows>) -> IO(Result<&1, &1, U32 & String, Vec<rows>>): MatA{m} = a VecA{xs} = x VecA{ys} = y IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<rows>>, B.Blas.sgemv(0, U32.from_nat(rows), U32.from_nat(cols), alpha, m, xs, beta, ys), r => Vec.wrap(rows, r))def Mat.gemv.try(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<cols>, beta: F32, y: Vec<rows>) -> IO(Vec<rows>): IO.try(Vec<rows>, Mat.gemv(rows, cols, alpha, a, x, beta, y))# y = alpha A^T x + beta ydef Mat.gemv_t(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<rows>, beta: F32, y: Vec<cols>) -> IO(Result<&1, &1, U32 & String, Vec<cols>>): MatA{m} = a VecA{xs} = x VecA{ys} = y IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<cols>>, B.Blas.sgemv(1, U32.from_nat(rows), U32.from_nat(cols), alpha, m, xs, beta, ys), r => Vec.wrap(cols, r))def Mat.gemv_t.try(+rows: Nat, +cols: Nat, alpha: F32, a: Mat<rows, cols>, x: Vec<rows>, beta: F32, y: Vec<cols>) -> IO(Vec<cols>): IO.try(Vec<cols>, Mat.gemv_t(rows, cols, alpha, a, x, beta, y))# A = alpha x y^T + Adef Mat.ger(+rows: Nat, +cols: Nat, alpha: F32, x: Vec<rows>, y: Vec<cols>, a: Mat<rows, cols>) -> IO(Result<&1, &1, U32 & String, Mat<rows, cols>>): VecA{xs} = x VecA{ys} = y MatA{m} = a IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<rows, cols>>, B.Blas.sger(U32.from_nat(rows), U32.from_nat(cols), alpha, xs, ys, m), r => Mat.wrap(rows, cols, r))def Mat.ger.try(+rows: Nat, +cols: Nat, alpha: F32, x: Vec<rows>, y: Vec<cols>, a: Mat<rows, cols>) -> IO(Mat<rows, cols>): IO.try(Mat<rows, cols>, Mat.ger(rows, cols, alpha, x, y, a))def Vec.dot(+n: Nat, x: Vec<n>, y: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>): VecA{xs} = x VecA{ys} = y B.Blas.sdot(U32.from_nat(n), xs, ys)def Vec.dot.try(+n: Nat, x: Vec<n>, y: Vec<n>) -> IO(F32): IO.try(F32, Vec.dot(n, x, y))def Vec.nrm2(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>): VecA{xs} = x B.Blas.snrm2(U32.from_nat(n), xs)def Vec.nrm2.try(+n: Nat, x: Vec<n>) -> IO(F32): IO.try(F32, Vec.nrm2(n, x))def Vec.asum(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, F32>): VecA{xs} = x B.Blas.sasum(U32.from_nat(n), xs)def Vec.asum.try(+n: Nat, x: Vec<n>) -> IO(F32): IO.try(F32, Vec.asum(n, x))def Vec.iamax(+n: Nat, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, U32>): VecA{xs} = x B.Blas.isamax(U32.from_nat(n), xs)def Vec.iamax.try(+n: Nat, x: Vec<n>) -> IO(U32): IO.try(U32, Vec.iamax(n, x))# y = alpha x + ydef Vec.axpy(+n: Nat, alpha: F32, x: Vec<n>, y: Vec<n>) -> IO(Result<&1, &1, U32 & String, Vec<n>>): VecA{xs} = x VecA{ys} = y IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<n>>, B.Blas.saxpy(U32.from_nat(n), alpha, xs, ys), r => Vec.wrap(n, r))def Vec.axpy.try(+n: Nat, alpha: F32, x: Vec<n>, y: Vec<n>) -> IO(Vec<n>): IO.try(Vec<n>, Vec.axpy(n, alpha, x, y))def Vec.scal(+n: Nat, alpha: F32, x: Vec<n>) -> IO(Result<&1, &1, U32 & String, Vec<n>>): VecA{xs} = x IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<n>>, B.Blas.sscal(U32.from_nat(n), alpha, xs), r => Vec.wrap(n, r))def Vec.scal.try(+n: Nat, alpha: F32, x: Vec<n>) -> IO(Vec<n>): IO.try(Vec<n>, Vec.scal(n, alpha, x))# LAPACK# ------# X = A^-1 B; Fail with the zero pivot as the code when A is singulardef Mat.solve(+n: Nat, +nrhs: Nat, a: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Result<&1, &1, U32 & String, Mat<n, nrhs>>): MatA{x} = a MatA{y} = b IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, nrhs>>, L.Lapack.sgesv(U32.from_nat(n), U32.from_nat(nrhs), x, y), r => Mat.wrap(n, nrhs, r))def Mat.solve.try(+n: Nat, +nrhs: Nat, a: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Mat<n, nrhs>): IO.try(Mat<n, nrhs>, Mat.solve(n, nrhs, a, b))# the lower Cholesky factor L of a symmetric positive definite Adef Mat.cholesky(+n: Nat, a: Mat<n, n>) -> IO(Result<&1, &1, U32 & String, Mat<n, n>>): MatA{x} = a IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, n>>, L.Lapack.spotrf(U32.from_nat(n), x), r => Mat.wrap(n, n, r))def Mat.cholesky.try(+n: Nat, a: Mat<n, n>) -> IO(Mat<n, n>): IO.try(Mat<n, n>, Mat.cholesky(n, a))# X = A^-1 B given L from Mat.choleskydef Mat.cholesky_solve(+n: Nat, +nrhs: Nat, l: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Result<&1, &1, U32 & String, Mat<n, nrhs>>): MatA{x} = l MatA{y} = b IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Mat<n, nrhs>>, L.Lapack.spotrs(U32.from_nat(n), U32.from_nat(nrhs), x, y), r => Mat.wrap(n, nrhs, r))def Mat.cholesky_solve.try(+n: Nat, +nrhs: Nat, l: Mat<n, n>, b: Mat<n, nrhs>) -> IO(Mat<n, nrhs>): IO.try(Mat<n, nrhs>, Mat.cholesky_solve(n, nrhs, l, b))# the singular values, largest firstdef Mat.svd_values(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)>>): MatA{x} = a IO.bind(Result<&1, &1, U32 & String, Array<F32>>, Result<&1, &1, U32 & String, Vec<Nat.min(m, n)>>, L.Lapack.sgesvd_s(U32.from_nat(m), U32.from_nat(n), x), r => Vec.wrap(Nat.min(m, n), r))def Mat.svd_values.try(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Vec<Nat.min(m, n)>): IO.try(Vec<Nat.min(m, n)>, Mat.svd_values(m, n, a))def Mat.svd.fin(-m: Nat, -n: Nat, s: Array<F32>, uv: Array<F32> & Array<F32>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>): match uv: case (u, vt): IO.pure(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, Done{(VecA{s}, (MatA{u}, MatA{vt}))})def Mat.svd.wrap(-m: Nat, -n: Nat, r: Result<&1, &1, U32 & String, Array<F32> & (Array<F32> & Array<F32>)>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>): match r: case Fail{e}: IO.pure(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, Fail{e}) case Done{v}: match v: case (s, uv): Mat.svd.fin(m, n, s, uv)# the thin SVD: (s, (u, vt)) with A = u diag(s) vtdef Mat.svd(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>): MatA{x} = a IO.bind(Result<&1, &1, U32 & String, Array<F32> & (Array<F32> & Array<F32>)>, Result<&1, &1, U32 & String, Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)>, L.Lapack.sgesvd(U32.from_nat(m), U32.from_nat(n), x), r => Mat.svd.wrap(m, n, r))def Mat.svd.try(+m: Nat, +n: Nat, a: Mat<m, n>) -> IO(Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>)): IO.try(Vec<Nat.min(m, n)> & (Mat<m, Nat.min(m, n)> & Mat<Nat.min(m, n), n>), Mat.svd(m, n, a))def Mat.eigh.wrap(-n: Nat, r: Result<&1, &1, U32 & String, Array<F32> & Array<F32>>) -> IO(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>): match r: case Fail{e}: IO.pure(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, Fail{e}) case Done{v}: match v: case (w, vs): IO.pure(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, Done{(VecA{w}, MatA{vs})})# the eigenvalues ascending and the eigenvectors as the rows of the Matdef Mat.eigh(+n: Nat, a: Mat<n, n>) -> IO(Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>): MatA{x} = a IO.bind(Result<&1, &1, U32 & String, Array<F32> & Array<F32>>, Result<&1, &1, U32 & String, Vec<n> & Mat<n, n>>, L.Lapack.ssyev(U32.from_nat(n), x), r => Mat.eigh.wrap(n, r))def Mat.eigh.try(+n: Nat, a: Mat<n, n>) -> IO(Vec<n> & Mat<n, n>): IO.try(Vec<n> & Mat<n, n>, Mat.eigh(n, a))