Interface for multi-dimensional arrays

Stream: lean4

[Tomas Skrivan]

I wanted to experiment with interface for multi-dimensional arrays. My main goal was to write down some basic formulas with the least amount of boilerplate as possible. I think, I have achieved a pretty good result, some basic linear algebra stuff:

def transpose (A : V2) : V2       := cmk λ i j => A[j,i]
def col (A : V2) (j : Fin 4) : V1 := cmk λ i => A[i,j]
def row (A : V2) (i : Fin 4) : V1 := cmk λ j => A[i,j]
def trace (A : V2) : ℝ            := ∑ i, A[i,i]

where Vn is a tensor of rank n(in this example every index has dimension 4) and cmk turns lambda into a tensor Vn. More complicated examples from general relativity:

def Γ₁ (g : V2) : V3 := cmk λ c a b => (D₂ g)[c,a,b] + (D₂ g)[c,b,a] - (D₂ g)[a,b,c]
def Γ₂ (g : V2) : V3 := cmk λ k i j => ∑ l, g[k,l]*(Γ₁ g)[l,i,j]
def R (g : V2) : V4 := cmk λ i j k l => let Γ : V3 := Γ₂ g
                                        (D₃ Γ)[l,i,k,j] + (D₃ Γ)[l,j,k,i] + ∑ p, (Γ[p,i,k] * Γ[l,j,p] - Γ[p,j,k] - Γ[l,i,p])
def 𝓡 (g : V2) : V2 := cmk λ i k => ∑ j, (R g)[i,j,k,j]
def SR (g : V2) : ℝ := ∑ i k, g[i,k] * (𝓡 g)[i,k]
def G (g : V2) : V2 := cmk λ i k => (𝓡 g)[i,k] - (SR g) * g[i,k]

At heart of this is NDArray class:

class NDArray (A : Type u) (T : Type v) (dims : Array Nat)  where
  get : A → NDArray.Index dims → T
  emk : (NDArray.Index dims → T) → A

That provides an interface to get an element of type T from (a : A) provided a multi-dimensional index and conversely can create the array given an element wise function.

I went for interface like design because A can store the multi-dimensional array in different ways. Column vs row major matrices, in CPU vs GPU memory, sparse/diagonal/block matrices etc. and I would like to have a common interface to all of these.

I would be happy for any comments, suggestions or critique.

[Tomas Skrivan]

The full code:

namespace NDArray

  def Index (dims : Array Nat) := (d : Fin dims.size) → Fin (dims[d])

end NDArray

--- Type A is a NDArray with densions dims and value type T
class NDArray (A : Type u) (T : Type v) (dims : Array Nat)  where
  get : A → NDArray.Index dims → T     -- get and element
  emk : (NDArray.Index dims → T) → A   -- elementa wise make

--- Automatically infering T and dims based on A
class NDArrayData (A : Type u) where
  T : Type v
  dims : Array Nat

-- Is this good idea?
@[reducible]
instance (A : Type u) (T : Type v) (dims : Array Nat) [NDArray A T dims] : NDArrayData A := ⟨T, dims⟩

attribute [reducible, inline] NDArrayData.T NDArrayData.dims

namespace NDArray

  namespace Index

    def toIndex1 (i1 : Fin n1) : Index #[n1]
    | Fin.mk 0 _ => i1

    def toIndex2 (i1 : Fin n1) (i2 : Fin n2) : Index #[n1, n2]
    | Fin.mk 0 _ => i1
    | Fin.mk 1 _ => i2

    def toIndex3 (i1 : Fin n1) (i2 : Fin n2) (i3 : Fin n3) : Index #[n1, n2, n3]
    | Fin.mk 0 _ => i1
    | Fin.mk 1 _ => i2
    | Fin.mk 2 _ => i3

    def toIndex4 (i1 : Fin n1) (i2 : Fin n2) (i3 : Fin n3) (i4 : Fin n4) : Index #[n1, n2, n3, n4]
    | Fin.mk 0 _ => i1
    | Fin.mk 1 _ => i2
    | Fin.mk 2 _ => i3
    | Fin.mk 3 _ => i4

    -- How to generalize?

  end Index

  @[reducible]
  abbrev scalarOf {A} (a : A) [NDArrayData A] := NDArrayData.T A

  @[reducible]
  abbrev dimsOf {A} (a : A) [NDArrayData A] := NDArrayData.dims A


  -- This can be turned into one macro once we have general toIndexₙ
  macro a:term noWs "[" i1:term "]" : term =>
    `(@NDArray.get _ (scalarOf $a) (dimsOf $a) _ $a (Index.toIndex1 $i1))

  macro a:term noWs "[" i1:term "," i2:term "]" : term =>
    `(@NDArray.get _ (scalarOf $a) (dimsOf $a) _ $a (Index.toIndex2 $i1 $i2))

  macro a:term noWs "[" i1:term "," i2:term "," i3:term "]" : term =>
    `(@NDArray.get _ (scalarOf $a) (dimsOf $a) _ $a (Index.toIndex3 $i1 $i2 $i3))

  macro a:term noWs "[" i1:term "," i2:term "," i3:term "," i4:term "]" : term =>
    `(@NDArray.get _ (scalarOf $a) (dimsOf $a) _ $a (Index.toIndex4 $i1 $i2 $i3 $i4))


  -- Make NDArray from an arbitrary type
  -- Mainly used to create an array from lambdas like (λ i j k => f i j k)
  section CustomMk
    class CustomMk (A : Type u) (α : Type w) where customMk : α → A

    variable {A : Type u} {T : Type v}

    instance [NDArray A T #[n]] : CustomMk A (Fin n → T) :=
             ⟨λ f => NDArray.emk (λ i : Index #[n] => f (i ⟨0, by simp[Array.size, List.length] done⟩))⟩
    instance [NDArray A T #[n1, n2]] : CustomMk A (Fin n1 → Fin n2 → T) :=
             ⟨λ f => NDArray.emk (λ i : Index #[n1, n2] => f (i ⟨0, by simp[Array.size, List.length] done⟩)
                                                             (i ⟨1, by simp[Array.size, List.length] done⟩))⟩
    instance [NDArray A T #[n1, n2, n3]] : CustomMk A (Fin n1 → Fin n2 → Fin n3 → T) :=
             ⟨λ f => NDArray.emk (λ i : Index #[n1, n2, n3] => f (i ⟨0, by simp[Array.size, List.length] done⟩)
                                                                 (i ⟨1, by simp[Array.size, List.length] done⟩)
                                                                 (i ⟨2, by simp[Array.size, List.length] done⟩))⟩
    instance [NDArray A T #[n1, n2, n3, n4]] : CustomMk A (Fin n1 → Fin n2 → Fin n3 → Fin n4 → T) :=
             ⟨λ f => NDArray.emk (λ i : Index #[n1, n2, n3, n4] => f (i ⟨0, by simp[Array.size, List.length] done⟩)
                                                                     (i ⟨1, by simp[Array.size, List.length] done⟩)
                                                                     (i ⟨2, by simp[Array.size, List.length] done⟩)
                                                                     (i ⟨3, by simp[Array.size, List.length] done⟩))⟩
    --- ... and so on ...

    def cmk [CustomMk A α] (a : α) : A := CustomMk.customMk a

  end CustomMk

end NDArray


-- Some auxiliary definitions
class Zero (α : Type u) where
  zero : α

instance instOfNatZero [Zero α] : OfNat α (nat_lit 0) where
  ofNat := Zero.zero

def sum {n α} [Zero α] [Add α] (f : Fin n → α) : α := do
  let mut r := 0
  for i in [0:n] do
    r := r + f ⟨i, sorry⟩
  r

macro "∑" xs:Lean.explicitBinders ", " b:term : term => Lean.expandExplicitBinders `sum xs b

section Test

    open NDArray

    constant ℝ : Type
    instance : Add ℝ := sorry
    instance : Mul ℝ := sorry
    instance : Sub ℝ := sorry
    instance : Zero ℝ := sorry
    constant V1 : Type
    constant V2 : Type
    constant V3 : Type
    constant V4 : Type
    instance : NDArray V1 ℝ #[4] := sorry
    instance : NDArray V2 ℝ #[4,4] := sorry
    instance : NDArray V3 ℝ #[4,4,4] := sorry
    instance : NDArray V4 ℝ #[4,4,4,4] := sorry

    def transpose (A : V2) : V2       := cmk λ i j => A[j,i]
    def col (A : V2) (j : Fin 4) : V1 := cmk λ i => A[i,j]
    def row (A : V2) (i : Fin 4) : V1 := cmk λ j => A[i,j]
    def trace (A : V2) : ℝ            := ∑ i, A[i,i]
    def mul (A B : V2) : V2           := cmk (λ i j => ∑ k, A[i,k]*B[k,j])

    variable [Inhabited V2] [Inhabited V3] [Inhabited V4]
    constant D₁ : V1 → V2
    constant D₂ : V2 → V3
    constant D₃ : V3 → V4

    -- General Relativity formulas
    -- https://en.wikipedia.org/wiki/List_of_formulas_in_Riemannian_geometry

    def Γ₁ (g : V2) : V3 := cmk λ c a b => (D₂ g)[c,a,b] + (D₂ g)[c,b,a] - (D₂ g)[a,b,c]
    def Γ₂ (g : V2) : V3 := cmk λ k i j => ∑ l, g[k,l]*(Γ₁ g)[l,i,j]
    def R (g : V2) : V4 := cmk λ i j k l => let Γ : V3 := Γ₂ g
                                            (D₃ Γ)[l,i,k,j] + (D₃ Γ)[l,j,k,i] + ∑ p, (Γ[p,i,k] * Γ[l,j,p] - Γ[p,j,k] - Γ[l,i,p])
    def 𝓡 (g : V2) : V2 := cmk λ i k => ∑ j, (R g)[i,j,k,j]
    def SR (g : V2) : ℝ := ∑ i k, g[i,k] * (𝓡 g)[i,k]
    def G (g : V2) : V2 := cmk λ i k => (𝓡 g)[i,k] - (SR g) * g[i,k]

end Test

[Tomas Skrivan]

One concrete question is how to merge function toIndexₙ with macros defining notation A[i,j,k]? Right, now I have to write separate function and macro for each A[i], A[i,j], A[i,j,k], ... Not a big deal, but it would be be nice to have one code handling any rank.

[Tomas Skrivan]

The thing I do not like is that cmk needs to know the output type. For example this works:

def transpose (A : V2) : V2 := cmk λ i j => A[j,i]

but this does not

def transpose (A : V2) := cmk λ i j => A[j,i]

As it is currently done, cmk needs to know the storage type for the output e.g. V2. In practice, this could be row major matrix or column major matrix or something completely different. Not sure what to do about it.

[Tomas Skrivan]

Also all those tests break when Vn are declared with variable and not with constant. Not yet sure what is going on there.

[Tomas Skrivan]

I have modified the definition of NDArray to

class NDArray (A : Type v → Array Nat → Type u) where
  get {T dims} : A T dims → NDArray.Index dims → T
  emk {T dims} : (NDArray.Index dims → T) → A T dims

and the code is a bit cleaner.

[Tomas Skrivan]

Full code:

-- Some auxiliary definitions
class Zero (α : Type u) where
  zero : α

instance instOfNatZero [Zero α] : OfNat α (nat_lit 0) where
  ofNat := Zero.zero

def sum {n α} [Zero α] [Add α] (f : Fin n → α) : α := do
  let mut r := 0
  for i in [0:n] do
    r := r + f ⟨i, sorry⟩
  r

macro "∑" xs:Lean.explicitBinders ", " b:term : term => Lean.expandExplicitBinders `sum xs b

namespace NDArray

  def Index (dims : Array Nat) := (d : Fin dims.size) → Fin (dims[d])

end NDArray

-- Do I want to have rank as explicit argument of `A` ??
--- Type A is a NDArray with densions dims and value type T
class NDArray (A : Type v → Array Nat → Type u) where
  get {T dims} : A T dims → NDArray.Index dims → T
  emk {T dims} : (NDArray.Index dims → T) → A T dims

namespace NDArray

  namespace Index
    def toIndex1 (i1 : Fin n1) : Index #[n1]
    | Fin.mk 0 _ => i1

    def toIndex2 (i1 : Fin n1) (i2 : Fin n2) : Index #[n1, n2]
    | Fin.mk 0 _ => i1
    | Fin.mk 1 _ => i2

    def toIndex3 (i1 : Fin n1) (i2 : Fin n2) (i3 : Fin n3) : Index #[n1, n2, n3]
    | Fin.mk 0 _ => i1
    | Fin.mk 1 _ => i2
    | Fin.mk 2 _ => i3

    def toIndex4 (i1 : Fin n1) (i2 : Fin n2) (i3 : Fin n3) (i4 : Fin n4) : Index #[n1, n2, n3, n4]
    | Fin.mk 0 _ => i1
    | Fin.mk 1 _ => i2
    | Fin.mk 2 _ => i3
    | Fin.mk 3 _ => i4

    -- How to generalize?

  end Index

  -- This can be turned into one macro once we have general toIndexₙ
  macro a:term noWs "[" i1:term "]" : term =>
    `(get $a (Index.toIndex1 $i1))

  macro a:term noWs "[" i1:term "," i2:term "]" : term =>
    `(get $a (Index.toIndex2 $i1 $i2))

  macro a:term noWs "[" i1:term "," i2:term "," i3:term "]" : term =>
    `(get $a (Index.toIndex3 $i1 $i2 $i3))

  macro a:term noWs "[" i1:term "," i2:term "," i3:term "," i4:term "]" : term =>
    `(get $a (Index.toIndex4 $i1 $i2 $i3 $i4))

  -- Make NDArray from an arbitrary type
  -- Mainly used to create an array from lambdas like (λ i j k => f i j k)
  section CustomMk
    class CustomMk (α : Type w) (A : Type v → Array Nat → Type u) (T dims) where customMk : α → A T dims

    variable {A T}
    instance [NDArray A] : CustomMk (Fin n → T) A T #[n] :=
             ⟨λ f => NDArray.emk (λ i : Index #[n] => f (i ⟨0, by simp[Array.size, List.length] done⟩))⟩
    instance [NDArray A] : CustomMk (Fin n1 → Fin n2 → T) A T #[n1,n2] :=
             ⟨λ f => NDArray.emk (λ i : Index #[n1, n2] => f (i ⟨0, by simp[Array.size, List.length] done⟩)
                                                             (i ⟨1, by simp[Array.size, List.length] done⟩))⟩
    instance [NDArray A] : CustomMk (Fin n1 → Fin n2 → Fin n3 → T) A T #[n1,n2,n3] :=
             ⟨λ f => NDArray.emk (λ i : Index #[n1, n2, n3] => f (i ⟨0, by simp[Array.size, List.length] done⟩)
                                                                 (i ⟨1, by simp[Array.size, List.length] done⟩)
                                                                 (i ⟨2, by simp[Array.size, List.length] done⟩))⟩
    instance [NDArray A] : CustomMk (Fin n1 → Fin n2 → Fin n3 → Fin n4 → T) A T #[n1,n2,n3,n4] :=
             ⟨λ f => NDArray.emk (λ i : Index #[n1, n2, n3, n4] => f (i ⟨0, by simp[Array.size, List.length] done⟩)
                                                                     (i ⟨1, by simp[Array.size, List.length] done⟩)
                                                                     (i ⟨2, by simp[Array.size, List.length] done⟩)
                                                                     (i ⟨3, by simp[Array.size, List.length] done⟩))⟩

    --- ... and so on ...

    def cmk {α A T dims} [CustomMk α A T dims] (a : α) : A T dims := CustomMk.customMk a

  end CustomMk

end NDArray

section Test

    open NDArray

    variable {ℝ : Type} [Add ℝ] [Mul ℝ] [Sub ℝ] [Zero ℝ]
    variable {V : Type → Array Nat → Type} [NDArray V]

    def transpose (A : V ℝ #[n,m]) : V ℝ #[m,n]  := cmk λ i j => A[j,i]
    def col (A : V ℝ #[n,m]) (j : Fin m) : V ℝ #[n] := cmk λ i => A[i,j]
    def row (A : V ℝ #[n,m]) (i : Fin n) : V ℝ #[m] := cmk λ j => A[i,j]
    def mul (A : V ℝ #[n,m]) (B : V ℝ #[m,k]) : V ℝ #[n,k] := cmk (λ i j => ∑ k, A[i,k]*B[k,j])

    variable [∀ dims, Inhabited (V ℝ dims)]
    constant D₂ : (V ℝ #[n,m]) → (V ℝ #[n,m,4])
    constant D₃ : (V ℝ #[n,m,k]) → (V ℝ #[n,m,k,4])

    -- General Relativity formulas
    -- https://en.wikipedia.org/wiki/List_of_formulas_in_Riemannian_geometry

    variable (g : V ℝ #[4,4])

    def Γ₁ : V ℝ #[4,4,4] := cmk λ c a b => (D₂ g)[c,a,b] + (D₂ g)[c,b,a] - (D₂ g)[a,b,c]
    def Γ₂ : V ℝ #[4,4,4] := cmk λ k i j => ∑ l, g[k,l]*(Γ₁ g)[l,i,j]
    def R  : V ℝ #[4,4,4,4] := cmk λ i j k l => let Γ := Γ₂ g
                                            (D₃ Γ)[l,i,k,j] + (D₃ Γ)[l,j,k,i] + ∑ p, (Γ[p,i,k] * Γ[l,j,p] - Γ[p,j,k] - Γ[l,i,p])
    def 𝓡  : V ℝ #[4,4] := cmk λ i k => ∑ j, (R g)[i,j,k,j]
    def SR : ℝ := ∑ i k, g[i,k] * (𝓡 g)[i,k]
    def G  : V ℝ #[4,4] := cmk λ i k => (𝓡 g)[i,k] - (SR g) * g[i,k]

end Test

[Kyle Miller]

Tomas Skrivan said:

One concrete question is how to merge function toIndexₙ with macros defining notation A[i,j,k]? Right, now I have to write separate function and macro for each A[i], A[i,j], A[i,j,k], ... Not a big deal, but it would be be nice to have one code handling any rank.

If you switched to A(i,j,k), you probably could implement this using a CoeFun instance for coercing arrays into functions with a tuple argument. You could also probably use CoeFun to make A i j k work.

[Kyle Miller]

You might also take a look at Dex, which has some interesting ideas. It has multidimensional array types using a function-like arrow v : n=>m=>(Fin 2)=>Float, and rather than limiting arrays to being indexed by elements of Fin _, they can be indexed by any finite enumerable type (that's what n and m are there). So, in this example, that type is for arrays that are basically indexed by n × m × Fin 2. To index into an array, you don't use the usual function application notation, but rather v.i.j.k

[Tomas Skrivan]

I see how to make A i j k work with CoeFun but I do not see how to make generic A(i,j,k), assuming that here A is turned to Fin n × Fin m × Fin k -> R.

[Tomas Skrivan]

Thanks for the reference to Dex, I haven't heard of it and will have a look.

[Kyle Miller]

For the having-multiple-storage-types problem, something that seems useful is being able to defer storage -- if you do A + B + C, ideally there'd be one outer loop that loads corresponding elements of the three arrays and adds them, rather than computing A + B, storing that, then adding C. One way to do this is to treat all arrays as being functions that never store anything. For example, you might implement transpose like this, where V2 is just Fin 4 -> Fin 4 -> R:

def transpose (A : V2) : V2 := λ i j => A j i

and then when you want to actually save an intermediate result, you can use a function like dense : V2 -> V2 that saves intermediate results in some underlying Array and immediately turns it back into a function (thus memoizing results). There's potential for accidental bad performance, though, if you forget to memoize results in key places. For that, it might be good to have some type like dV2 for a specifically dense vectors, and some coercions between V2 and dV2. Then you could write, for example,

let b : dV2 := transpose a

[Tomas Skrivan]

Currently I'm restricting myself to Fin n because Lean 4 does not have fintype and right now I have different priorities then that.

[Kyle Miller]

The reason you might do this is for typechecking reasons, since it helps prevent you using indexes in the wrong spots when you write more generic functions. You also probably wouldn't want Fintype anyway, but rather a type with a specific equivalence to Fin n. I understand it not being high priority, though.

[Tomas Skrivan]

I understand it not being high priority, though.

I'm sure that someone else will do it for mathlib4 :) So I will just wait and get my part of the code ready.

[Tomas Skrivan]

Kyle Miller said:

For the having-multiple-storage-types problem, something that seems useful is being able to defer storage -- if you do A + B + C, ideally there'd be one outer loop that loads corresponding elements of the three arrays and adds them, rather than computing A + B, storing that, then adding C. One way to do this is to treat all arrays as being functions that never store anything. For example, you might implement transpose like this, where V2 is just Fin 4 -> Fin 4 -> R:

def transpose (A : V2) : V2 := λ i j => A j i

and then when you want to actually save an intermediate result, you can use a function like dense : V2 -> V2 that saves intermediate results in some underlying Array and immediately turns it back into a function (thus memoizing results). There's potential for accidental bad performance, though, if you forget to memoize results in key places. For that, it might be good to have some type like dV2 for a specifically dense vectors, and some coercions between V2 and dV2. Then you could write, for example,

let b : dV2 := transpose a

My idea to this problem is to take similar approach as Halide where you write the code in the most straight forward way and only then you decide how it should be evaluated.

I believe tactic mode can be used for that. Write the the code in a very simple way first and then apply bunch of rw tactics, or more specialized ones, to transform it into a faster implementation. The advantage is that you have a guarantee the code is doing what you indented.

[Kyle Miller]

Tomas Skrivan said:

I see how to make A i j k work with CoeFun but I do not see how to make generic A(i,j,k), assuming that here A is turned to Fin n × Fin m × Fin k -> R.

Part of it would be making a function that converts Array Nat into a cartesian product of types, which I guess could actually replace your Index function, so then the implementation of CoeFun would be NDArray.get.

[Tomas Skrivan]

Part of it would be making a function that converts Array Nat into a cartesian product of types

That is the problem, I have no clue how to do that.

[Kyle Miller]

Here's some code from an experiment from a while back. I don't know if it works (it's old, and I copy-pasted the relevant parts), but it has the equivalence stuff I was talking about.

def Function.leftInverse (g : β → α) (f : α → β) : Prop :=
    ∀ x, g (f x) = x
def Function.rightInverse (g : β → α) (f : α → β) : Prop :=
    Function.leftInverse f g

structure Equiv (α : Sort u) (β : Sort v) where
    toFun : α → β
    invFun : β → α
    leftInv : Function.leftInverse invFun toFun
    rightInv : Function.rightInverse invFun toFun

infix:25 " ≃ " => Equiv

def Equiv.symm (f : α ≃ β) : β ≃ α where
    toFun := f.invFun
    invFun := f.toFun
    leftInv := f.rightInv
    rightInv := f.leftInv

/-- An equivalence "is" a function. -/
instance (α : Type u) (β : Type v) : CoeFun (Equiv α β) (λ _ => α → β) where
    coe := Equiv.toFun

class Enumerable (α : Type u) where
    card : Nat
    enum : α ≃ Fin card

instance : Enumerable (Fin n) where
    card := n
    enum := {
        toFun := id
        invFun := id
        leftInv := λ _ => rfl
        rightInv := λ _ => rfl
    }

section CartesianProduct

theorem cartEncodeProp {i j m n : Nat} (hi : i < m) (hj : j < n) : i * n + j < m * n := by
    cases m with
    | zero => exfalso; exact Nat.notLtZero _ hi
    | succ m => {
        rw Nat.succMul;
        exact Nat.ltOfLeOfLt (Nat.addLeAddRight (Nat.mulLeMulRight _ (Nat.leOfLtSucc hi)) _) (Nat.addLtAddLeft hj _)
    }

def cartDecode {n m : Nat} (k : Fin (n * m)) : Fin n × Fin m :=
    let ⟨k, h⟩ := k
    (
        ⟨k / m, sorry⟩,
        ⟨k % m, Nat.modLt _ (by { cases m; exfalso; rw Nat.mulZero at h; exact Nat.notLtZero _ h; apply Nat.succPos})⟩
    )

instance [Enumerable α] [Enumerable β] : Enumerable (α × β) where
    card := Enumerable.card α * Enumerable.card β
    enum := {
        toFun := λ (a, b) =>
            let ⟨i, hi⟩ := Enumerable.enum a
            let ⟨j, hj⟩ := Enumerable.enum b
            ⟨i * Enumerable.card β + j, cartEncodeProp hi hj⟩
        invFun := λ n =>
            let (i, j) := cartDecode n
            (Enumerable.enum.symm i, Enumerable.enum.symm j)
        leftInv := sorry
        rightInv := sorry
    }

end CartesianProduct

def Enumerable.listOf.aux (α : Type u) [Enumerable α] : Nat -> Nat -> List α
| lo, 0 => []
| lo, (left+1) =>
    if h : lo < Enumerable.card α then
        Enumerable.enum.symm ⟨lo, h⟩ :: aux α (lo + 1) left
    else [] -- Shouldn't happen, but makes the definition easy.

/-- Create a list of every term in the Enumerable type in order. -/
def Enumerable.listOf (α : Type u) [Enumerable α] : List α :=
    Enumerable.listOf.aux α 0 (Enumerable.card α)

def Vec.sum [Enumerable ι] [Add α] [OfNat α Nat.zero] (v : Vec ι α) : α := do
    let mut s : α := 0
    for i in Enumerable.listOf ι do
        s := s + v[i]
    return s

/-- A function-backed vector -/
structure Vec (ι : Type u) (α : Type v) where
    toFun : ι → α

macro "vec" xs:Lean.explicitBinders " => " b:term : term => Lean.expandExplicitBinders `Vec.mk xs b

/-- support `v[i]` notation. -/
@[inline] def Vec.getOp (self : Vec ι α) (idx : ι) : α := self.toFun idx

instance [Add α] : Add (Vec ι α) where
    add v w := vec i => v[i] + w[i]
instance [Sub α] : Sub (Vec ι α) where
    sub v w := vec i => v[i] - w[i]
instance [Mul α] : Mul (Vec ι α) where
    mul v w := vec i => v[i] * w[i]
instance [Div α] : Div (Vec ι α) where
    div v w := vec i => v[i] / w[i]

namespace Vec

def push (v : Vec ι (Vec κ α)) : Vec (ι × κ) α := vec p => v[p.1][p.2]
def pop (v : Vec (ι × κ) α) : Vec ι (Vec κ α) := vec i j => v[(i, j)]
def reindex (v : Vec ι α) (f : κ → ι) : Vec κ α := vec i => v[f i]

instance : Monad (Vec ι) where
    pure x := vec i => x
    map f v := vec i => f v[i]
    seq f v := vec i => f[i] v[i]
    bind v f := vec i => (f v[i])[i] -- diagonal (is this actually a monad law?)

instance [OfNat α n] : OfNat (Vec ι α) n where
    ofNat := vec i => OfNat.ofNat n

def transpose (v : Vec ι (Vec κ α)) : Vec κ (Vec ι α) := vec j i => v[i][j]

end Vec

def Vec.sum [Enumerable ι] [Add α] [OfNat α Nat.zero] (v : Vec ι α) : α := do
    let mut s : α := 0
    for i in Enumerable.listOf ι do
        s := s + v[i]
    return s

[Tomas Skrivan]

Well one way would be:

def toFinProd : List Nat → Type
| List.nil => Unit
| n :: ns => Fin n × toFinProd ns

but that turns [n,m] into FIn n × (Fin m × Unit). At some point, I was playing around with types like that and I found them super painful to use.

[Kyle Miller]

If you special-case one-element lists, it gives you a nicer type:

def toCart : List Nat → Type
| [] => Unit
| [n] => Fin n
| (n::ns) => Fin n × toCart ns

[Kyle Miller]

I'm not sure, but you might need to make it @[reducible] or something to make sure Lean unfolds it for typeclass inference. It would be nicer if toCart [3,4,5] actually looked like Fin 3 × Fin 4 × Fin 5.

[Tomas Skrivan]

Ohh even getting the notation A i j k will involve something like toCart. You need to turn d : List Nat into Fin d[0] -> Fin d[1] -> ... -. Fin d[n-1] -> R.

[Tomas Skrivan]

Btw. thanks for the code but one big difference is that I definitely do not want to go the route of having matrices defined as Vec ι (Vec κ α). What about diagonal matrices, block matrices, sparse matrices etc.

[Kyle Miller]

The part of the code I didn't copy, since I didn't really get it to work right, was having things like DenseVec that can be used as a Vec.

[Kyle Miller]

Vec is just the ideal mathematical type, and then there would be different concrete storage types.

[Tomas Skrivan]

Ohh I see, I wasn't thinking about the code in the right way. I have to keep staring at it :)

[Kyle Miller]

This might not matter for what you're doing, but something that I was sort of going for was to be able to do APL/J-style array computations, where you can think of your multidimensional array as being a multidimensional array of multidimensional arrays and broadcast function applications over each of the arrays.

[Tomas Skrivan]

At some point I would like to do something like that too, so getting something working would be nice.

[Tomas Skrivan]

How would you specify the storage type in your approach?

[Tomas Skrivan]

With your approach, you can easily do multidimensional arrays as v : Vec (Fin 3 × Fin 4 × Fin 5) R then v[(i,j,k)] works and turning it into v[i,j,k] is just a simple macro.

[Kyle Miller]

I had imagined making types that can be coerced to Vec, so DenseVec might be a structure containing an Array of all the elements, or a DiagonalMatrix might be a structure containing an Array of just the diagonal elements.

[Tomas Skrivan]

A mix of yours and my approach would be:

class NDArray (A : Type u) (ι : Type v) (α : Type w) where
  toFun : A → ι → α     -- get and element

[Tomas Skrivan]

Kyle Miller said:

I had imagined making types that can be coerced to Vec, so DenseVec might be a structure containing an Array of all the elements, or a DiagonalMatrix might be a structure containing an Array of just the diagonal elements.

Makes sense, so I'm curious which part you didn't manage to get working?

[Kyle Miller]

Here's another experiment for Lean 3 (though nothing about storage types): https://gist.github.com/kmill/5eb83cb5fd102e4242c2c8a53224ba13

As an illustration, two kinds of matrix-vector products depending on whether the indices have been merged into a cartesian product:

def matvec {r s α : Type*} [fintype s] [comm_ring α] (A : r => s => α) (v : s => α) : r => α :=
for i, vsum for j, A i j * v j

def matvec' {r s α : Type*} [fintype s] [comm_ring α] (A : r × s => α) (v : s => α) : r => α :=
for i, vsum for j, A (i, j) * v j

There are some operations for going back and forth between these representations:

def vec.push {ι₁ ι₂ α : Type*} (v : ι₁ => ι₂ => α) : ι₁ × ι₂ => α
def vec.pop {ι₁ ι₂ α : Type*} (v : ι₁ × ι₂ => α) : ι₁ => ι₂ => α

Or using applicative-style for broadcasting addition:

def addvec' {r α : Type u} [add_monoid α] (v w : r => α) : r => α :=
(+) <$> v <*> w

It also has an extremely fiddly notation (which would be easier in Lean 4) for slicing: u⟦(),j⟧ and u⟦i,()⟧for a u : a=>b=>c.

[Kyle Miller]

Tomas Skrivan said:

I'm curious which part you didn't manage to get working?

Unfortunately, I don't really remember... I think it was something to do with the ergonomics of using a DenseVec in an expression, since Lean had no reason to try coercing it to a Vec for an operation like addition. It looks like I had my own coercion class

class HasVec (β : Type u) (ι : outParam $ Type v) (α : outParam $ Type w) where
    toVec : β → Vec ι α

and then you'd write something like 2 * toVec d + toVec d'

[Tomas Skrivan]

Wouldn't that get solved by using the

class NDArray (A : Type u) (ι : Type v) (α : Type w) where
  toFun : A → ι → α     -- get and element

An equivalent to Vec would be the instance

instance : NDArray (ι → α) ι α := ⟨λ f => f⟩

Then you are not relaying on coercion but on typeclasses. Addition would be

instance [NDArray A ι α] [NDArray B ι α] : HAdd A B (ι → α) := ...

and you do not need the custom coercion toVec.

[Tomas Skrivan]

Broadcasting addition

instance [NDArray A I B] [NDArray B J R] : HAdd A B (I → J → R) := ...

[Kyle Miller]

I imagine that should work. (I had tried something like that with HAdd, but this was shortly after the public release of Lean 4, and there may have potentially been some typeclass inference bugs I ran into that have since been patched.)

[Tomas Skrivan]

I will play around with it a bit, scared of associativity of × though.

[Tomas Skrivan]

I have done a quick experiment and this works:

variable (f : Fin 3 × Fin 2 ↦ Nat) (g : Fin 2 × Fin 5 ↦ Nat) (h : Fin 3 × Fin 5 ↦ Nat)
variable (r s : Nat)

#check f[1,0]
#check λ i j => f[i,j]
#check f + f
#check r * f * g + s * h

where α ↦ β is just α → β but it is a hint that element notation f[i,j] can be used.

[Tomas Skrivan]

Full code

--- Container `C` with index set `ι` and element type `α`
class Cont (C : Type u) (ι : Type v) (α : Type w) where
  toFun : C → ι → α

--- Automatically infering T and dims based on A
class ContData (C : Type u) where
  indexType : Type v
  valueType : Type w

-- Is this good idea?
@[reducible]
instance (C : Type u) (ι : Type v) (α : Type w) [Cont C ι α] : ContData C := ⟨ι, α⟩

attribute [reducible, inline] ContData.indexType ContData.valueType


namespace Cont

  -- Function that should be interpreted as a container
  def ContFun (α β) := α → β
  infix:34 " ↦ " => ContFun

  def toCont (f : α → β) : α ↦ β := f
  instance (ι : Type v) (α : Type w) : Cont (ι ↦ α) ι α := ⟨λ f => f⟩

  @[reducible]
  abbrev indexOf {C} (c : C) [ContData C] := ContData.indexType C

  @[reducible]
  abbrev valueOf {C} (c : C) [ContData C] := ContData.valueType C

  @[reducible]
  abbrev get {C} [ContData C] [Cont C (ContData.indexType C) (ContData.valueType C)] (c : C) := @toFun _ (indexOf c) (valueOf c) _ c

  macro c:term noWs "[" i1:term"]" : term =>
    `(get $c $i1)

  macro c:term noWs "[" i1:term "," i2:term "]" : term =>
    `(get $c ($i1,$i2))

  macro c:term noWs "[" i1:term "," i2:term "," i3:term "]" : term =>
    `(get $c ($i1,$i2,$i3))

  macro c:term noWs "[" i1:term "," i2:term "," i3:term "," i4:term "]" : term =>
    `(get $c ($i1,$i2,$i3,$i4))


  section Arithmetics

     section ElementWise
       variable {ι : Type v}
       variable {C : Type u} {α : Type w} [Cont C ι α]
       variable {C' : Type u'} {α' : Type w'} [Cont C' ι α']

       instance [HAdd α α' β] : HAdd C C' (ι ↦ β) := ⟨λ c c' => λ i => c[i] + c'[i]⟩

     end ElementWise

     section Mul
       variable {ι κ μ : Type v}
       variable {C : Type u} {α : Type w} [Cont C (ι×κ) α]
       variable {C' : Type u'} {α' : Type w'} [Cont C' (κ×μ) α']

       -- Just a prototype because I do not have a `Fintype` defined
       def sum [Add β] : (α → β) → β := sorry

       instance [HMul α α' β] [Add β] : HMul C C' (ι × μ ↦ β) := ⟨λ c c' (i,j) => sum λ k => c[i,k] * c'[k,j]⟩

     end Mul

     section Broadcasting

       variable {ι : Type v}
       variable {C : Type u} {α : Type w} [Cont C ι α]

       instance [HMul α β γ] : HMul C β (ι ↦ γ) := ⟨λ a b i => a[i]*b⟩
       instance [HMul β α γ] : HMul β C (ι ↦ γ) := ⟨λ b a i => b*a[i]⟩


     end Broadcasting

  end Arithmetics

end Cont


variable (f : Fin 3 × Fin 2 ↦ Nat) (g : Fin 2 × Fin 5 ↦ Nat) (h : Fin 3 × Fin 5 ↦ Nat)
variable (r s : Nat)

#check f[1,0]
#check λ i j => f[i,j]
#check f + f
#check r * f * g + s * h