Skip to content

Commit

Permalink
some doodling with pooling layers
Browse files Browse the repository at this point in the history
  • Loading branch information
lecopivo committed Nov 22, 2023
1 parent db9a9f9 commit 6f77b13
Show file tree
Hide file tree
Showing 2 changed files with 278 additions and 25 deletions.
272 changes: 272 additions & 0 deletions SciLean/Modules/ML/Doodle.lean
Original file line number Diff line number Diff line change
@@ -0,0 +1,272 @@
import SciLean
import SciLean.Core.Meta.GenerateRevCDeriv'

namespace SciLean

variable
{K : Type} [RealScalar K]
{W : Type} [Vec K W]

set_default_scalar K

variable {α β κ ι : Type} [Index.{0,0,0} α] [Index.{0,0,0} β] [Index.{0,0,0} κ] [Index.{0,0,0} ι] [PlainDataType K] [PlainDataType R]

variable (κ)
def denseLazy (weights : κ → ι → K) (bias : κ → K) (x : ι → K) (j : κ) : K :=
∑ i, weights j i * x i + bias j
variable {κ}

#generate_revCDeriv' denseLazy weights bias x | j
prop_by unfold denseLazy; fprop
trans_by
unfold denseLazy
autodiff

variable (κ)
def dense (weights : DataArrayN K (κ×ι)) (bias : K^κ) (x : K^ι) : K^κ :=
-- ⊞ j => ∑ i, weights[(j,i)] * x[i] + bias[j]
⊞ j => denseLazy κ (fun j i => weights[(j,i)]) (fun j => bias[j]) (fun i => x[i]) j
variable {κ}

#generate_revCDeriv' dense weights bias x
prop_by unfold dense; fprop
trans_by unfold dense; autodiff



#eval 0

#exit
set_option synthInstance.maxSize 20000


#check denseLazy.arg_weightsbiasx_j.revCDeriv
#check dense.arg_weightsbiasx.revCDeriv

variable (x : K^(Idx 20))

set_option profiler true in

#check (revCDeriv K fun (w,w',w'',w''',w'''',b,b',b'',b''',b'''') =>
x |> dense (Idx 5) w' b'
|> dense (Idx 10) w b
|> dense (Idx 20) w'' b''
|> dense (Idx 20) w''' b'''
|> dense (Idx 30) w'''' b'''')
rewrite_by
ftrans only
ftrans only
ftrans only
-- lsimp (config := {zeta:=false, singlePass:=true}) only


example (a : Nat) (b : Nat)
: (b,a,a,0) + (b,a,a,a,a,a,a,a,a,b,b) + 0 = (b+b,a+a,a+a,a,a,a,a,a,a,b,b) := by simp

example (a : K)
: (a,0) + (a,a,a,a,a,a,a,a,a,a,a) + 0 = (a+a,a,a,a,a,a,a,a,a,a,a) := by simp

example (a : K ^ Idx 10) (b : K ^ Idx 20)
: (a,0) + (a,a,a,a,a,a,a,a,a,b,b) + 0 = (a+a,a,a,a,a,a,a,a,a,b,b) := by simp


set_option trace.Meta.Tactic.simp.discharge true in
example (a : K ^ Idx 10) (b : K ^ Idx 20)
: (b,a,a,0) + (b,a,a,a,a,a,a,a,a,b,b) + 0 = (b+b,a+a,a+a,a,a,a,a,a,a,b,b) := by simp


set_option trace.Meta.Tactic.simp.discharge true in
example (a : K ^ Idx 10)
: (a,a,a,0) + (a,a,a,a,a,a,a,a,a,a,a) = (a+a,a+a,a+a,a,a,a,a,a,a,a,a) := by simp


variable (a : Nat)
#check
(a +
(a + d
let x := a + a
x))
rewrite_by
lsimp (config := {zeta := false, singlePass := true})

#exit

set_option trace.Meta.Tactic.ftrans.step true
set_option trace.Meta.Tactic.ftrans.theorems true
set_option trace.Meta.Tactic.fprop.discharge true
set_option trace.Meta.Tactic.simp.discharge true
set_option trace.Meta.Tactic.simp.congr true
set_option trace.Meta.Tactic.simp.rewrite true
set_option trace.Meta.Tactic.simp.unify true
example [SemiInnerProductSpace K W]
: <∂ (fun (x : ((W × DataArrayN K (κ × ι)) × K ^ κ) × K ^ ι) (j : κ) =>
denseLazy κ (fun (j : κ) (i : ι) => x.fst.fst.snd[(j, i)]) (fun (j : κ) => x.fst.snd[j]) (fun (i : ι) => x.snd[i]) j)
=
fun _ => 0 :=
by
conv =>
lhs
ftrans only


#check SciLean.denseLazy.arg_weightsbiasx_j.revCDeriv_rule_def
#check SciLean.denseLazy.arg_weightsbiasx_j.revCDeriv_rule
#check SciLean.denseLazy.arg_weightsbiasx_j.HasAdjDiff_rule

#exit
variable {W : Type _} [SemiInnerProductSpace K W]

example (x : W → DataArrayN K ι) (hx : ∀ i, HasAdjDiff K (fun w => (x w)[i]))
: HasAdjDiff K x := by fprop

example (x : W → DataArrayN K ι) (i : ι) (hx : HasAdjDiff K x)
: HasAdjDiff K fun w => (x w)[i] := by fprop

example (x : W → DataArrayN K ι) (hx : HasAdjDiff K x)
: HasAdjDiff K fun w i => (x w)[i] := by fprop


-- def foo : Float → DataArrayN Float (Idx 10) := sorry

-- -- set_option maxHeartbeats 10000

-- -- set_option trace.Meta.isDefEq true in
-- open Lean Meta Qq in
-- #eval show MetaM Unit from do

-- let X : Q(Type) := q(Float → DataArrayN Float (Idx 10))
-- withLocalDecl `x default X fun x => do
-- let x : Q($X) := x
-- let HX := q(IsDifferentiable Float $x)
-- withLocalDecl `hx default HX fun hx => do

-- let H := q(IsDifferentiable Float fun w => ⊞ i => ($x w)[i])
-- let h ← mkFreshExprMVar H
-- IO.println (← isDefEq hx h)



-- set_option trace.Meta.isDefEq true in
-- open Lean Meta Qq in
-- #eval show MetaM Unit from do

-- let X : Q(Type) := q(Float → DataArrayN Float (Idx 10))
-- withLocalDecl `x default X fun x => do
-- let x : Q($X) := x
-- let HX := q(IsDifferentiable Float $x)
-- withLocalDecl `hx default HX fun hx => do

-- let H := q(IsDifferentiable Float fun w => ⊞ i => ($x w)[i])
-- let h ← mkFreshExprMVar H
-- IO.println (← isDefEq hx h)

-- set_option maxHeartbeats 50000


-- set_option trace.Meta.Tactic.fprop.step true in
-- set_option trace.Meta.Tactic.fprop.unify true in
-- set_option trace.Meta.Tactic.fprop.discharge true in

example
(x : Float → DataArrayN Float (Idx 10)) (hx : IsDifferentiable Float x)
: IsDifferentiable Float (fun w => ⊞ i => (x w)[i]) :=
by
fprop


variable (x : Float → DataArrayN Float (Idx 10)) (hx : ∀ i : Idx 10, HasAdjDiff Float (fun w => (x w)[i]))

set_default_scalar Float

#check <∂ w, x w


example
: (<∂ w, ∑ i, (x w)[i])
=
fun w =>
let xdx := <∂ sorry :=
by
sorry

set_option trace.Meta.Tactic.simp.discharge true in
set_option trace.Meta.Tactic.simp.rewrite true in
#check
<∂ w, ∑ i, (x w)[i]
rewrite_by
simp (config := {zeta := false}) only [SciLean.EnumType.sum.arg_f.revCDeriv_rule _ sorry]
simp (config := {zeta := false}) only [SciLean.revCDeriv.pi_rule _ _ sorry]


simp (config := {zeta := false}) only
ftrans only



@[fprop]
theorem dense.arg_weightsbiasx.IsDifferentiable_rule
(weights : W → DataArrayN K (κ×ι)) (bias : W → DataArrayN K κ) (x : W → DataArrayN K ι)
(hweights : IsDifferentiable K weights) (hbias : IsDifferentiable K bias) (hx : IsDifferentiable K x)
: IsDifferentiable K fun w => dense κ (weights w) (bias w) (x w) :=
by
unfold dense; unfold denseLazy
fprop


@[ftrans]
theorem dense.arg_weightsbiasx.fwdCDeriv_rule
(weights : W → DataArrayN K (κ×ι)) (bias : W → DataArrayN K κ) (x : W → DataArrayN K ι)
(hweights : IsDifferentiable K weights) (hbias : IsDifferentiable K bias) (hx : IsDifferentiable K x)
: (fwdCDeriv K fun w => dense κ (weights w) (bias w) (x w) )
=
((fwdCDeriv K fun w => dense κ (weights w) (bias w) (x w))
rewrite_by unfold dense; unfold denseLazy; autodiff) :=
by
unfold dense; unfold denseLazy
conv => lhs; autodiff


set_option trace.Meta.Tactic.ftrans.step true in
@[ftrans]
theorem dense.arg_weightsbiasx.revCDeriv_rule {W : Type} [SemiInnerProductSpace K W]
(weights : W → DataArrayN K (κ×ι)) (bias : W → DataArrayN K κ) (x : W → DataArrayN K ι)
(hweights : HasAdjDiff K weights) (hbias : HasAdjDiff K bias) (hx : HasAdjDiff K x)
: (revCDeriv K fun w => dense κ (weights w) (bias w) (x w) )
=
((revCDeriv K fun w => dense κ (weights w) (bias w) (x w))
rewrite_by unfold dense; unfold denseLazy; autodiff; autodiff) :=
by
sorry


#exit
section denseDerivTest
variable (weights : DataArrayN R (κ×ι)) (bias : DataArrayN R κ) (x : DataArrayN R ι)

#check
∇ x, dense κ weights bias x
rewrite_by
unfold dense; symdiff

#check
∇ bias, dense κ weights bias x
rewrite_by
unfold dense; symdiff

#check
∇ weights, dense κ weights bias x
rewrite_by
unfold dense; symdiff

end denseDerivTest


structure Decomposition (X X₁ X₂ : Type) where
split : X → X₁ × X₂
merge : X₁ → X₂ → X


example : (g : G) [Curry Xs Y] (dec : Decomposition Xs Xs₁ Xs₂) (f : F) [Curry F (Y×Xs₁) Z]
: revCDeriv fun xs =>
let y := uncurry' g xs
uncurry' f (y, dec.split xs)
31 changes: 6 additions & 25 deletions SciLean/Modules/ML/Pool.lean
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import SciLean.Data.Prod
import SciLean.Core.Meta.GenerateRevCDeriv'

import SciLean.Core.FloatAsReal

import SciLean.Core.Function

namespace SciLean

Expand All @@ -15,29 +15,6 @@ set_default_scalar R

variable {κ ι κ'} [Index κ] [Index κ'] [Index ι] [PlainDataType R]

def _root_.Function.reduce {ι α} [Index ι] [Inhabited α] (f : ι → α) (op : α → α → α) : α := Id.run do
let n := Index.size ι
if 0 = n then
return default
else
let mut a ← f (fromIdx ⟨0, sorry_proof⟩)
for i in [1:n.toNat] do
a ← op a (← f (fromIdx ⟨i.toUSize, sorry_proof⟩))
return a

theorem _root_.Function.reducte.arg_f.revCDeriv {ι K X} [Index ι] [IsROrC K] [SemiInnerProductSpace K X]
(f : ι → X) (dop : X → X×X) : X × (X → (ι→X)) := Id.run do
let n := Index.size ι
if 0 = n then
return (default, 0)
let mut a : Array X := Array.mkEmpty n.toNat
let mut da : Array X := Array.mkEmpty n.toNat

sorry

#eval (fun i : Idx 5 => i.1).reduce (·+·)

#check ForInStep.yield

example
{ι X : Type} [Index ι]
Expand Down Expand Up @@ -105,6 +82,10 @@ def poolLazy
(op : R → R → R)
(x : ι → R)
(j : κ) : R :=
Index.joinl (fun j' : κ' => x (indexSplit.symm (j,j'))) op
Function.reduce (fun j' : κ' => x (indexSplit.symm (j,j'))) op

variable {κ}

#generate_revCDeriv' poolLazy x | j
prop_by unfold poolLazy; fprop
trans_by unfold poolLazy; ftrans

0 comments on commit 6f77b13

Please sign in to comment.