synthesizing prop guards

Stream: new members

[samuela]

Hi folks! I'm new to Lean, and was curious to cut my teeth on modeling NumPy-style broadcasting in Lean at the type level. Array shapes are represented as List Nat and I have taken the approach of using a Prop to encode that two shapes are Broadcastable s1 s2. However, I'm running into two issues:

1. I'm finding that a Broadcastable s1 s2 implicit parameter is not being synthesized as I would expect in HAdd, and
2. No matter what I do, I can't get lean to synthesize the Broadcastable s1 s2 argument in the application of HAdd.add

I'm curious: Why aren't these arguments being synthesized by Lean? I have found that sometimes arguments are synthesized correctly, but sometimes not. (more details in code) More broadly, is this an appropriate design for this problem?

Ideally, I would like an approach that used maximum automation to fill in these instances. They are tedious and ought to be easy work for something like grind. I have heard of Decidable, but I'm not sure how that works or if it would be applicable here?

MWE code

[Aaron Liu]

You need to make it a typeclass

[Aaron Liu]

That's the square bracket system

[Robin Arnez]

There are basically two types of synthesization: one for type classes and one for regular holes/implicits. Type class synthesization happens when you use a function with an instance implicit parameter ([MyClass ...]) and uses variables of the form MyClass ... from the local context and any instance declarations. The other kind of synthesization is inferring values from typing rules. For example when you have l : List Nat and do List.length l, then it can infer the implicit parameter α from context because it expects l : List ?α and gets l : List Nat. On the other hand, this won't take in variables from the local context. For example you wouldn't want this to randomly assign using the Nat from the context:

def mkVector {n : Nat} : Vector Nat n := Vector.replicate n 0

/--
error: don't know how to synthesize implicit argument `n`
  @mkVector (?m.5 n)
context:
n : Nat
⊢ Nat
---
error: Failed to infer type of definition `test`
-/
#guard_msgs in
def test (n : Nat) := mkVector

For example, you could've also meant @mkVector 10 or @mkVector (n + 1) etc. so instead of using the n from the local context, it instead simply tells you that it doesn't have enough information to infer the parameter n.

[samuela]

Ok, thank you! So IIUC typeclass synthesis, aka [Foo] notation, will do the whole rigmarole of synthesis to chain any number of type classes together to get any instance that satisfactorily fills a type class hole (a la haskell, etc) and implicit arguments, aka {_ : Foo} notation, will only synthesize from other variables in a function definition as in the List Nat example? So {}-synthesis will not utilize things in context at site of the function application

[Robin Arnez]

To be perfectly precise, {_ : Foo} arguments might also be filled in through other means, especially the expected type but that's the gist. Also worth noting is that the behavior of implicit arguments is equivalent to the behavior of the _/(_) term

[samuela]

ok i'm trying out a type class approach: code but i'm getting an error:

cannot find synthesization order for instance @instBroadcastableTrans with type
  ∀ {s₁ s₂ s₃ : Shape} [Broadcastable s₁ s₂] [Broadcastable s₂ s₃], Broadcastable s₁ s₃
all remaining arguments have metavariables:
  Broadcastable s₁ ?s₂
  Broadcastable ?s₂ s₃

on what feels like an inoffensive instance. am i missing something? should type class resolution be able to handle this?

[Riccardo Brasca]

The problem is that Lean cannot guess s₂ only looking at the goal.

[Riccardo Brasca]

So this cannot be an instance: otherwise anytime Lean sees a goal of type Broadcastable foo bar it has to try Broadcastable foo X and Broadcastable X bar for all possible X, and it gets crazy.

[samuela]

no dice with this route either

[samuela]

mm i see, ok that makes sense

[Riccardo Brasca]

Note that it can be a def without any problem.

[samuela]

hmm how would the def look? would that be compatible with getting broadcast type checking on the shapes when doing eg x + y?

[Riccardo Brasca]

Just write def instead of instance.

[Riccardo Brasca]

Note that I am not really reading the code, just explaining the general error.

[Robin Arnez]

Here's a basic version based on the original MWE:

namespace Foo

/-- A tensor "shape" is a list of natural dimensions, e.g. `[2, 3]` for a 2x3 matrix. -/
abbrev Shape := List Nat

/-- Compute the NumPy-style broadcasted shape, assuming as a precondition that the inputs are safe to broadcast (aligning from the right). -/
def broadcastRev : Shape → Shape → Shape
  | [], ys => ys
  | xs, [] => xs
  | x :: xs, y :: ys => Nat.max x y :: broadcastRev xs ys

/-- Broadcast result shape for two shapes, aligning from the right. -/
def broadcast (s₁ s₂ : Shape) : Shape :=
  (broadcastRev s₁.reverse s₂.reverse).reverse

-- l₁.reverse ++ l₂
abbrev revAuxReducible (l₁ l₂ : List α) : List α :=
  match l₁ with
  | [] => l₂
  | a :: l => revAuxReducible l (a :: l₂)

abbrev revReducible (l : List α) : List α := revAuxReducible l []

set_option checkBinderAnnotations false in
class inductive BroadcastableRev : Shape → Shape → Prop where
  | nil_left (ys) : BroadcastableRev [] ys
  | nil_right (xs) : BroadcastableRev xs []
  | step_equal {x xs ys}
      [BroadcastableRev xs ys] :
      BroadcastableRev (x :: xs) (x :: ys)
  | step_left1 {xs ys y}
      [BroadcastableRev xs ys] :
      BroadcastableRev (1 :: xs) (y :: ys)
  | step_right1 {xs ys x}
      [BroadcastableRev xs ys] :
      BroadcastableRev (x :: xs) (1 :: ys)

open BroadcastableRev in
attribute [instance] nil_left nil_right step_equal step_left1 step_right1

/-- Right-aligned broadcastability on the original (non-reversed) shapes. -/
abbrev Broadcastable (s₁ s₂ : Shape) : Prop :=
  BroadcastableRev (revReducible s₁) (revReducible s₂)

/--
An array carrying only its element type `α` and its static shape `shape`. We do
not model any runtime storage for now.
-/
structure Array (α : Type) (shape : Shape) : Type where
  /-- Dummy payload for now. Just playing around with types. -/
  tracer : Unit := ()

/--
Elementwise addition with NumPy-style broadcasting at the type level.

Requires a proof that the two shapes are broadcastable; the result shape
is the computed `broadcast s₁ s₂`.
-/
def addBroadcast {α : Type} {s₁ s₂ : Shape} [Broadcastable s₁ s₂] :
  Array α s₁ → Array α s₂ → Array α (broadcast s₁ s₂) :=
  fun _ _ => ⟨()⟩

/-- Support the generalized `+` notation via `HAdd`. -/
instance {α : Type} {s₁ s₂ : Shape} [Broadcastable s₁ s₂] :
    HAdd (Array α s₁) (Array α s₂) (Array α (broadcast s₁ s₂)) where
  hAdd := addBroadcast

#check (let x : Array Nat [2, 3] := ⟨()⟩
        let y : Array Nat [2, 3] := ⟨()⟩
        x + y)

(you can open it in Lean 4 Web using the button in the top-right corner of the code block)

[samuela]

woah very cool! checking it out now

[samuela]

ok, out of curiosity i tried making all of broadcast be abbrev, so that it would expand those things in the output types, but it doesn't seem to make a difference, eg Lean still reports the types of things as Array Nat (broadcast [2, 3] [1, 1, 3]) instead of Array Nat [1, 2, 3] as I would expect. I'm curious why that is? is there a way to get it to elaborate those things? (code)

[Robin Arnez]

abbrev does not literally mean "abbreviation" in the sense of macros. It just means that in the majority of cases, tactics and term elaborators can look through the abbreviation as if it wasn't there (even if it is there!). You can achieve your idea by instead "encoding" broadcast in Broadcastable. The way to do is so called outParams which get filled in during instance synthesis:

namespace Foo

abbrev Shape := List Nat

set_option checkBinderAnnotations false in
class inductive BroadcastableSameSize : Shape → Shape → outParam Shape → Prop where
  | refl (xs) : BroadcastableSameSize xs xs xs
  | step_equal [BroadcastableSameSize xs ys zs] : BroadcastableSameSize (x :: xs) (x :: ys) (x :: zs)
  | step_left1 [BroadcastableSameSize xs ys zs] : BroadcastableSameSize (1 :: xs) (y :: ys) (y :: zs)
  | step_right1 [BroadcastableSameSize xs ys zs] : BroadcastableSameSize (x :: xs) (1 :: ys) (x :: zs)

set_option checkBinderAnnotations false in
class inductive BroadcastableLeft : Shape → Shape → outParam Shape → Prop where
  | of_same_size [BroadcastableSameSize xs ys zs] : BroadcastableLeft xs ys zs
  | skip_left [BroadcastableLeft xs ys zs] : BroadcastableLeft (x :: xs) ys (x :: zs)

set_option checkBinderAnnotations false in
class inductive BroadcastableRight : Shape → Shape → outParam Shape → Prop where
  | of_same_size [BroadcastableSameSize xs ys zs] : BroadcastableRight xs ys zs
  | skip_right [BroadcastableRight xs ys zs] : BroadcastableRight xs (y :: ys) (y :: zs)

set_option checkBinderAnnotations false in
class inductive Broadcastable : Shape → Shape → outParam Shape → Prop where
  | of_same_size [BroadcastableSameSize xs ys zs] : Broadcastable xs ys zs
  | skip_left [BroadcastableLeft xs ys zs] : Broadcastable (x :: xs) ys (x :: zs)
  | skip_right [BroadcastableRight xs ys zs] : Broadcastable xs (y :: ys) (y :: zs)

open Broadcastable BroadcastableSameSize in
attribute [instance high] refl of_same_size BroadcastableLeft.of_same_size BroadcastableRight.of_same_size
open Broadcastable BroadcastableSameSize in
attribute [instance] step_equal step_left1 step_right1 skip_left skip_right BroadcastableLeft.skip_left BroadcastableRight.skip_right

structure Array (α : Type) (shape : Shape) : Type where
  tracer : Unit := ()

def addBroadcast {α : Type} {s₁ s₂ s₃ : Shape} [Broadcastable s₁ s₂ s₃] :
    Array α s₁ → Array α s₂ → Array α s₃ :=
  fun _ _ => ⟨()⟩

instance [Broadcastable s₁ s₂ s₃] : HAdd (Array α s₁) (Array α s₂) (Array α s₃) where
  hAdd := addBroadcast

#check (let x : Array Nat [2, 1] := ⟨()⟩
        let y : Array Nat [1, 3] := ⟨()⟩
        x + y)

[Robin Arnez]

(there's probably a better way to express the condition that doesn't require 4 typeclasses...)

[samuela]

ooh nice! ok i was not aware of outParam. i will check it out