sourcedemos/lora-training/lora-training.xtl
1⍝!/usr/bin/env xetal
2⍝│ LoRA training: where the numbers of A and B come from. A frozen 8 x
3⍝│ 8 layer W; examples (inputs with the outputs wanted) made by a
4⍝│ hidden shift of it, W + T; an adapter that starts as LoRA does, B
5⍝│ at 0 and A seeded. One step of training as a pipeline, stage by
6⍝│ stage: the prediction, the error (which is the loss's gradient at
7⍝│ the output), the gradients of A and B, the update; checked against
8⍝│ the library's own step, W untouched. Then training by Adam with a
9⍝│ record: the loss on the examples trained on and on examples held
10⍝│ out, every few steps, and the adapter at each record, so a page
11⍝│ can step back without training again.
12
13"e:" u̲se< "Equations"
14"flow:" u̲se< "Flow"
15"lo:" u̲se< "LoRA"
16"p:" u̲se< "Plot"
17
18⍝━━ Data
19
20⍝│ The settings (a preset): the rank, alpha (the strength s is alpha /
21⍝│ r), the learning rate, how many of the 12 examples are trained on
22⍝│ (the last 4 are always held out), and the record: chunks of every
23⍝│ steps.
24r ← 4
25alpha ← 4.0
26lr ← 0.05
27n ← 8
28chunks ← 30
29every ← 10
30⍝│ The frozen layer W, by a fixed ripple (the lab's), and the hidden
31⍝│ shift T (8 x 8, of rank 4) that the examples come from: the layer
32⍝│ W + T is what the adapter is to imitate, without ever seeing T.
33i ← f̲loat r̲ange 64
34W ← 8 8 r̲eshape s̲in 0.7 × i × i
35p ← f̲loat r̲ange 32
36T ← (8 4 r̲eshape 0.5 × s̲in 1.3 × p × p) '+ '× i̲nner 4 8 r̲eshape 0.5 × c̲os 0.6 × p × p
37⍝│ 12 inputs of 8 numbers, one per column, by a ripple; the output
38⍝│ wanted for each is the shifted layer's. The first n columns are
39⍝│ trained on, the last 4 held out: never trained on, only measured.
40c ← f̲loat r̲ange 96
41X ← 8 12 r̲eshape c̲os 0.4 × c + 0.05 × c × c
42Y ← (W + T) '+ '× i̲nner X
43Xt ← n t̲ake₂ X
44Yt ← n t̲ake₂ Y
45Xh ← -4 t̲ake₂ X
46Yh ← -4 t̲ake₂ Y
47⍝│ The adapter as LoRA starts it: B all 0, so the correction is 0 and
48⍝│ the layer gives exactly W's output; A seeded (fixed, no chance).
49s ← alpha ˡᵒs̲caling r
50q ← r − 1
51A ← 1 ˡᵒs̲eeded r c̲at 8
52B ← (8 c̲at r) r̲eshape 0.0
53a0 ← (W, A, B, s)
54⍝│ The task's tables, for the page: W, T, the inputs and the outputs
55⍝│ wanted (by rows).
56r̲avel W
57r̲avel T
58r̲avel X
59r̲avel Y
60
61⍝━━ One step of training
62
63⍝│ One step of gradient descent on the n examples trained on, a stage
64⍝│ at a time. The prediction p (the layer as it is: W x plus the
65⍝│ correction, 0 at the start); the error g = p - y, which is also the
66⍝│ loss's gradient at the output (the loss is half the sum of the
67⍝│ squared errors, lo:l_oss); the error carried back through B, b =
68⍝│ B' g; the gradients of A and B, each a sum over the examples; the
69⍝│ update, A and B moved against their gradients by the learning
70⍝│ rate. W is not in the update: the step gives it back as it was,
71⍝│ as the adapter after the step, (W2, A3, B3, s2).
72step ← @ ᶠˡᵒʷp̲ipe< c#"
73 Xt ▹x
74▷ "h[k,m] | r n" ᵉe̲q< "sum j:=0..7 | A[k,j] * x[j,m]" ▹h
75▷ "u[i,m] | 8 n" ᵉe̲q< "sum k:=0..q | B[i,k] * h[k,m]" ▹u
76▷ "p[i,m] | 8 n" ᵉe̲q< "(sum j:=0..7 | W[i,j] * x[j,m]) + s * u[i,m]" ▹p
77▷ "g[i,m] | 8 n" ᵉe̲q< "p[i,m] - Yt[i,m]" ▹g
78▷ "b[k,m] | r n" ᵉe̲q< "sum i:=0..7 | B[i,k] * g[i,m]" ▹b
79▷ "dA[k,j] | r 8" ᵉe̲q< "sum m:=0..n-1 | s * b[k,m] * x[j,m]" ▹dA
80▷ "dB[i,k] | 8 r" ᵉe̲q< "sum m:=0..n-1 | s * g[i,m] * h[k,m]" ▹dB
81▷ "A2[k,j] | r 8" ᵉe̲q< "A[k,j] - lr * dA[k,j]" ▹A2
82▷ "B2[i,k] | 8 r" ᵉe̲q< "B[i,k] - lr * dB[i,k]" ▹B2
83▷ (W, A2, B2, s) ▹(W2, A3, B3, s2)"#
84⍝│ The loss before the step and after it, on the examples trained on,
85⍝│ per example (lo:l_oss gives the sum over the examples; dividing by
86⍝│ their number makes the loss on n examples comparable with the loss
87⍝│ on the 4 held out).
88(a0 ˡᵒl̲oss (Xt, Yt)) ÷ f̲loat n
89(step ˡᵒl̲oss (Xt, Yt)) ÷ f̲loat n
90⍝│ The library's own step (lo:s_gd) from the same start: how far the
91⍝│ pipeline's A and B are from its (0 to rounding), and how far its W
92⍝│ is from W (exactly 0: the step does not touch it).
93(Ws, As, Bs, _) ← lr ˡᵒs̲gd (a0, (Xt, Yt))
94(W2, A2, B2, _) ← step
95'm̲ax r̲/ (r̲avel a̲bs A2 − As) c̲at r̲avel a̲bs B2 − Bs
96'm̲ax r̲/ r̲avel a̲bs (W2 − W) c̲at Ws − W
97
98⍝━━ Training
99
100⍝│ Training by Adam (lo:a_dam), chunks times every steps, with a
101⍝│ record after each chunk: the loss per example on the examples
102⍝│ trained on, the same on the held-out ones, and A, B and the
103⍝│ correction s B A as they are then (flattened, one record after
104⍝│ another), the start recorded first.
105(_, lossT, lossH, recA, recB, recD) ← chunks '{ (z, lt, lh, ra, rb, rd) →
106 z2 ← every '{ e → lr ˡᵒa̲dam (e, (Xt, Yt)) } p̲ower z
107 (a, _, _) ← z2
108 (_, a2, b2, _) ← a
109 (z2, lt c̲at (a ˡᵒl̲oss (Xt, Yt)) ÷ f̲loat n, lh c̲at (a ˡᵒl̲oss (Xh, Yh)) ÷ 4.0, ra c̲at r̲avel a2, rb c̲at r̲avel b2, rd c̲at r̲avel ˡᵒd̲elta (a2, b2, s))
110} p̲ower (ˡᵒa̲damStart a0, (a0 ˡᵒl̲oss (Xt, Yt)) ÷ f̲loat n, (a0 ˡᵒl̲oss (Xh, Yh)) ÷ 4.0, r̲avel A, r̲avel B, r̲avel ˡᵒd̲elta (A, B, s))
111lossT
112lossH
113recA
114recB
115recD
116⍝│ The two losses as one chart, a record a point.
117_ ← ("Training" "record" "loss per example" "trained on" "held out") ᵖc̲hart! '{ k → k s̲elect (2 c̲at chunks + 1) r̲eshape lossT c̲at lossH } m̲ap r̲ange 2
ᶠˡᵒʷp̲ipe< expands to
({ @ →
x ← Xt
h ← "h[k,m] | r n" ᵉe̲q< "sum j:=0..7 | A[k,j] * x[j,m]"
u ← "u[i,m] | 8 n" ᵉe̲q< "sum k:=0..q | B[i,k] * h[k,m]"
p ← "p[i,m] | 8 n" ᵉe̲q< "(sum j:=0..7 | W[i,j] * x[j,m]) + s * u[i,m]"
g ← "g[i,m] | 8 n" ᵉe̲q< "p[i,m] - Yt[i,m]"
b ← "b[k,m] | r n" ᵉe̲q< "sum i:=0..7 | B[i,k] * g[i,m]"
dA ← "dA[k,j] | r 8" ᵉe̲q< "sum m:=0..n-1 | s * b[k,m] * x[j,m]"
dB ← "dB[i,k] | 8 r" ᵉe̲q< "sum m:=0..n-1 | s * g[i,m] * h[k,m]"
A2 ← "A2[k,j] | r 8" ᵉe̲q< "A[k,j] - lr * dA[k,j]"
B2 ← "B2[i,k] | 8 r" ᵉe̲q< "B[i,k] - lr * dB[i,k]"
(W2, A3, B3, s2) ← (W, A2, B2, s)
(W2, A3, B3, s2)
} @)ᵉe̲q< expands to
(({ @ →
g3:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g2:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g1:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g4:grid ← r c̲at n
g5:g̲rid ← { k → g4:grid r̲eshape (o̲ffsets k s̲elect g4:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g4:grid) r̲eshape 0 }
g4:grid r̲eshape (({ @ →
g6:grid ← r c̲at n c̲at ({ v → (0 ≤ v) ? v◆ "bad-equation" ⎕S̲IGNAL "sum's range runs low..high: 0..7" } (1 + 7 − 0))
g7:g̲rid ← { k → g6:grid r̲eshape (o̲ffsets k s̲elect g6:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g6:grid) r̲eshape 0 }
'+ r̲/₃ g6:grid r̲eshape ((
((g1:r̲ank (A, "A", 2, "A[k,j]")) g3:a̲t ((
((A, "A", 2, "j") g2:c̲heck (0 + g7:g̲rid 3)) + (2 s̲elect s̲hape A) ×
((A, "A", 1, "k") g2:c̲heck (0 + g7:g̲rid 1))))) ×
((g1:r̲ank (x, "x", 2, "x[j,m]")) g3:a̲t ((
((x, "x", 2, "m") g2:c̲heck (0 + g7:g̲rid 2)) + (2 s̲elect s̲hape x) ×
((x, "x", 1, "j") g2:c̲heck (0 + g7:g̲rid 3)))))))
} @))
} @))ᵉe̲q< expands to
(({ @ →
g10:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g9:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g8:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g11:grid ← 8 c̲at n
g12:g̲rid ← { k → g11:grid r̲eshape (o̲ffsets k s̲elect g11:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g11:grid) r̲eshape 0 }
g11:grid r̲eshape (({ @ →
g13:grid ← 8 c̲at n c̲at ({ v → (0 ≤ v) ? v◆ "bad-equation" ⎕S̲IGNAL "sum's range runs low..high: 0..q" } (1 + q − 0))
g14:g̲rid ← { k → g13:grid r̲eshape (o̲ffsets k s̲elect g13:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g13:grid) r̲eshape 0 }
'+ r̲/₃ g13:grid r̲eshape ((
((g8:r̲ank (B, "B", 2, "B[i,k]")) g10:a̲t ((
((B, "B", 2, "k") g9:c̲heck (0 + g14:g̲rid 3)) + (2 s̲elect s̲hape B) ×
((B, "B", 1, "i") g9:c̲heck (0 + g14:g̲rid 1))))) ×
((g8:r̲ank (h, "h", 2, "h[k,m]")) g10:a̲t ((
((h, "h", 2, "m") g9:c̲heck (0 + g14:g̲rid 2)) + (2 s̲elect s̲hape h) ×
((h, "h", 1, "k") g9:c̲heck (0 + g14:g̲rid 3)))))))
} @))
} @))ᵉe̲q< expands to
(({ @ →
g17:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g16:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g15:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g18:grid ← 8 c̲at n
g19:g̲rid ← { k → g18:grid r̲eshape (o̲ffsets k s̲elect g18:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g18:grid) r̲eshape 0 }
g18:grid r̲eshape ((({ @ →
g20:grid ← 8 c̲at n c̲at ({ v → (0 ≤ v) ? v◆ "bad-equation" ⎕S̲IGNAL "sum's range runs low..high: 0..7" } (1 + 7 − 0))
g21:g̲rid ← { k → g20:grid r̲eshape (o̲ffsets k s̲elect g20:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g20:grid) r̲eshape 0 }
'+ r̲/₃ g20:grid r̲eshape ((
((g15:r̲ank (W, "W", 2, "W[i,j]")) g17:a̲t ((
((W, "W", 2, "j") g16:c̲heck (0 + g21:g̲rid 3)) + (2 s̲elect s̲hape W) ×
((W, "W", 1, "i") g16:c̲heck (0 + g21:g̲rid 1))))) ×
((g15:r̲ank (x, "x", 2, "x[j,m]")) g17:a̲t ((
((x, "x", 2, "m") g16:c̲heck (0 + g21:g̲rid 2)) + (2 s̲elect s̲hape x) ×
((x, "x", 1, "j") g16:c̲heck (0 + g21:g̲rid 3)))))))
} @) + (s ×
((g15:r̲ank (u, "u", 2, "u[i,m]")) g17:a̲t ((
((u, "u", 2, "m") g16:c̲heck (0 + g19:g̲rid 2)) + (2 s̲elect s̲hape u) ×
((u, "u", 1, "i") g16:c̲heck (0 + g19:g̲rid 1))))))))
} @))ᵉe̲q< expands to
(({ @ →
g24:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g23:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g22:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g25:grid ← 8 c̲at n
g26:g̲rid ← { k → g25:grid r̲eshape (o̲ffsets k s̲elect g25:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g25:grid) r̲eshape 0 }
g25:grid r̲eshape ((
((g22:r̲ank (p, "p", 2, "p[i,m]")) g24:a̲t ((
((p, "p", 2, "m") g23:c̲heck (0 + g26:g̲rid 2)) + (2 s̲elect s̲hape p) ×
((p, "p", 1, "i") g23:c̲heck (0 + g26:g̲rid 1))))) −
((g22:r̲ank (Yt, "Yt", 2, "Yt[i,m]")) g24:a̲t ((
((Yt, "Yt", 2, "m") g23:c̲heck (0 + g26:g̲rid 2)) + (2 s̲elect s̲hape Yt) ×
((Yt, "Yt", 1, "i") g23:c̲heck (0 + g26:g̲rid 1)))))))
} @))ᵉe̲q< expands to
(({ @ →
g29:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g28:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g27:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g30:grid ← r c̲at n
g31:g̲rid ← { k → g30:grid r̲eshape (o̲ffsets k s̲elect g30:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g30:grid) r̲eshape 0 }
g30:grid r̲eshape (({ @ →
g32:grid ← r c̲at n c̲at ({ v → (0 ≤ v) ? v◆ "bad-equation" ⎕S̲IGNAL "sum's range runs low..high: 0..7" } (1 + 7 − 0))
g33:g̲rid ← { k → g32:grid r̲eshape (o̲ffsets k s̲elect g32:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g32:grid) r̲eshape 0 }
'+ r̲/₃ g32:grid r̲eshape ((
((g27:r̲ank (B, "B", 2, "B[i,k]")) g29:a̲t ((
((B, "B", 2, "k") g28:c̲heck (0 + g33:g̲rid 1)) + (2 s̲elect s̲hape B) ×
((B, "B", 1, "i") g28:c̲heck (0 + g33:g̲rid 3))))) ×
((g27:r̲ank (g, "g", 2, "g[i,m]")) g29:a̲t ((
((g, "g", 2, "m") g28:c̲heck (0 + g33:g̲rid 2)) + (2 s̲elect s̲hape g) ×
((g, "g", 1, "i") g28:c̲heck (0 + g33:g̲rid 3)))))))
} @))
} @))ᵉe̲q< expands to
(({ @ →
g36:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g35:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g34:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g37:grid ← r c̲at 8
g38:g̲rid ← { k → g37:grid r̲eshape (o̲ffsets k s̲elect g37:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g37:grid) r̲eshape 0 }
g37:grid r̲eshape (({ @ →
g39:grid ← r c̲at 8 c̲at ({ v → (0 ≤ v) ? v◆ "bad-equation" ⎕S̲IGNAL "sum's range runs low..high: 0..n-1" } (1 + (n − 1) − 0))
g40:g̲rid ← { k → g39:grid r̲eshape (o̲ffsets k s̲elect g39:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g39:grid) r̲eshape 0 }
'+ r̲/₃ g39:grid r̲eshape (((s ×
((g34:r̲ank (b, "b", 2, "b[k,m]")) g36:a̲t ((
((b, "b", 2, "m") g35:c̲heck (0 + g40:g̲rid 3)) + (2 s̲elect s̲hape b) ×
((b, "b", 1, "k") g35:c̲heck (0 + g40:g̲rid 1)))))) ×
((g34:r̲ank (x, "x", 2, "x[j,m]")) g36:a̲t ((
((x, "x", 2, "m") g35:c̲heck (0 + g40:g̲rid 3)) + (2 s̲elect s̲hape x) ×
((x, "x", 1, "j") g35:c̲heck (0 + g40:g̲rid 2)))))))
} @))
} @))ᵉe̲q< expands to
(({ @ →
g43:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g42:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g41:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g44:grid ← 8 c̲at r
g45:g̲rid ← { k → g44:grid r̲eshape (o̲ffsets k s̲elect g44:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g44:grid) r̲eshape 0 }
g44:grid r̲eshape (({ @ →
g46:grid ← 8 c̲at r c̲at ({ v → (0 ≤ v) ? v◆ "bad-equation" ⎕S̲IGNAL "sum's range runs low..high: 0..n-1" } (1 + (n − 1) − 0))
g47:g̲rid ← { k → g46:grid r̲eshape (o̲ffsets k s̲elect g46:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g46:grid) r̲eshape 0 }
'+ r̲/₃ g46:grid r̲eshape (((s ×
((g41:r̲ank (g, "g", 2, "g[i,m]")) g43:a̲t ((
((g, "g", 2, "m") g42:c̲heck (0 + g47:g̲rid 3)) + (2 s̲elect s̲hape g) ×
((g, "g", 1, "i") g42:c̲heck (0 + g47:g̲rid 1)))))) ×
((g41:r̲ank (h, "h", 2, "h[k,m]")) g43:a̲t ((
((h, "h", 2, "m") g42:c̲heck (0 + g47:g̲rid 3)) + (2 s̲elect s̲hape h) ×
((h, "h", 1, "k") g42:c̲heck (0 + g47:g̲rid 2)))))))
} @))
} @))ᵉe̲q< expands to
(({ @ →
g50:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g49:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g48:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g51:grid ← r c̲at 8
g52:g̲rid ← { k → g51:grid r̲eshape (o̲ffsets k s̲elect g51:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g51:grid) r̲eshape 0 }
g51:grid r̲eshape ((
((g48:r̲ank (A, "A", 2, "A[k,j]")) g50:a̲t ((
((A, "A", 2, "j") g49:c̲heck (0 + g52:g̲rid 2)) + (2 s̲elect s̲hape A) ×
((A, "A", 1, "k") g49:c̲heck (0 + g52:g̲rid 1))))) − (lr ×
((g48:r̲ank (dA, "dA", 2, "dA[k,j]")) g50:a̲t ((
((dA, "dA", 2, "j") g49:c̲heck (0 + g52:g̲rid 2)) + (2 s̲elect s̲hape dA) ×
((dA, "dA", 1, "k") g49:c̲heck (0 + g52:g̲rid 1))))))))
} @))ᵉe̲q< expands to
(({ @ →
g55:a̲t ← { A i → (1 + i) s̲elect r̲avel A }
g54:c̲heck ← { (A, name, k, from) e →
n ← k s̲elect s̲hape A
bad ← w̲here n̲ot r̲avel (0 ≤ e) ∧ e < n
0 = t̲ally bad ? e
"index" ⎕S̲IGNAL name c̲at " axis " c̲at (f̲ormat k) c̲at ": index " c̲at (f̲ormat (f̲irst bad) s̲elect r̲avel e) c̲at " is outside 0.." c̲at (f̲ormat n − 1) c̲at " (from " c̲at from c̲at ")"
}
g53:r̲ank ← { (A, name, q, ref) →
(t̲ally s̲hape A) = q ? A
"rank" ⎕S̲IGNAL ref c̲at " indexes " c̲at (f̲ormat q) c̲at " axes, but " c̲at name c̲at " has rank " c̲at f̲ormat t̲ally s̲hape A
}
g56:grid ← 8 c̲at r
g57:g̲rid ← { k → g56:grid r̲eshape (o̲ffsets k s̲elect g56:grid) 'l̲eft t̲able ('× r̲/ k d̲rop g56:grid) r̲eshape 0 }
g56:grid r̲eshape ((
((g53:r̲ank (B, "B", 2, "B[i,k]")) g55:a̲t ((
((B, "B", 2, "k") g54:c̲heck (0 + g57:g̲rid 2)) + (2 s̲elect s̲hape B) ×
((B, "B", 1, "i") g54:c̲heck (0 + g57:g̲rid 1))))) − (lr ×
((g53:r̲ank (dB, "dB", 2, "dB[i,k]")) g55:a̲t ((
((dB, "dB", 2, "k") g54:c̲heck (0 + g57:g̲rid 2)) + (2 s̲elect s̲hape dB) ×
((dB, "dB", 1, "i") g54:c̲heck (0 + g57:g̲rid 1))))))))
} @))