1	# A convolution followed by an activation, as a training layer: forward keeps
2	# what backward reads, backward gives the gradients in a fixed FP32 order.
3	import Base
4	import ./math_activation.bend as Activations
5	import ./tensor.bend as Tensors
6	import ./tensor_f32.bend as F
7	import ./tensor_convolution_gradient.bend as Gradient
8	
9	def Tensor() -> Type:
10	  Tensors.Tensor<F32>
11	
12	type Settings is Data:
13	  Settings{stride: U32,padding: U32,activation: Activations.Activation}
14	
15	def settings(stride: U32,padding: U32,activation: Activations.Activation) -> Settings:
16	  Settings{stride,padding,activation}
17	
18	# The layer input and the convolution result before the activation. One
19	# backward call uses it.
20	def Saved() -> Type:
21	  Tensor() & Tensor()
22	
23	def forward_activated(x: Tensor(),weights: Tensor(),bias: Tensor(),pair: Tensor() & Tensor()) -> Tensor() & (Saved() & (Tensor() & Tensor())):
24	  (output,z) = pair
25	  (output,((x,z),(weights,bias)))
26	
27	def forward_copied(settings: Settings,xs: Tensor() & Tensor(),ws: Tensor() & Tensor(),bs: Tensor() & Tensor()) -> Tensor() & (Saved() & (Tensor() & Tensor())):
28	  Settings{stride,padding,+activation} = settings
29	  (x,x_kept) = xs
30	  (weights,weights_kept) = ws
31	  (bias,bias_kept) = bs
32	  forward_activated(x_kept,weights_kept,bias_kept,F.activate_forward(F.conv2d(x,weights,bias,stride,padding,Activations.identity()),activation))
33	
34	# x [N,C,H,W] or [C,H,W], weights [O,C,k,k] and bias [O]:
35	# (the activated output, (what backward reads, (weights, bias))).
36	def forward(x: Tensor(),weights: Tensor(),bias: Tensor(),settings: Settings) -> Tensor() & (Saved() & (Tensor() & Tensor())):
37	  forward_copied(settings,F.copy(x),F.copy(weights),F.copy(bias))
38	
39	def kernel(shape: List<&2,U32>) -> U32:
40	  match shape:
41	    case [outputs,channels,kernel,width]: kernel
42	    case shape: 0
43	
44	# The gradients of the parameters with what the input gradient reads:
45	# (weight gradient, (bias gradient, (dz, (weights, the shape of x)))).
46	def Parameters() -> Type:
47	  Tensor() & (Tensor() & (Tensor() & (Tensor() & List<&2,U32>)))
48	
49	def biased(weight_gradient: Tensor(),weights: Tensor(),shape: List<&2,U32>,pair: Tensor() & Tensor()) -> Parameters():
50	  (bias_gradient,dz) = pair
51	  (weight_gradient,(bias_gradient,(dz,(weights,shape))))
52	
53	def measured(weight_gradient: Tensor(),weights: Tensor(),dz: Tensor(),input: Tensor() & List<&2,U32>) -> Parameters():
54	  (x,shape) = input
55	  biased(weight_gradient,weights,shape,Gradient.bias(dz))
56	
57	def weighted(weights: Tensor(),result: Tensor() & (Tensor() & Tensor())) -> Parameters():
58	  (weight_gradient,(x,dz)) = result
59	  measured(weight_gradient,weights,dz,Tensors.shape(~F32,x))
60	
61	def shaped(stride: U32,padding: U32,x: Tensor(),dz: Tensor(),measured: Tensor() & List<&2,U32>) -> Parameters():
62	  (weights,shape) = measured
63	  weighted(weights,Gradient.weight(x,dz,kernel(shape),stride,padding))
64	
65	# dz = dy * activation'(z), then the gradients of the weights and the bias.
66	def parameters(dy: Tensor(),saved: Saved(),weights: Tensor(),settings: Settings) -> Parameters():
67	  (x,z) = saved
68	  Settings{stride,padding,+activation} = settings
69	  shaped(stride,padding,x,F.activate_backward(dy,z,activation),Tensors.shape(~F32,weights))
70	
71	def finished(weight_gradient: Tensor(),bias_gradient: Tensor(),pair: Tensor() & Tensor()) -> Tensor() & (Tensor() & (Tensor() & Tensor())):
72	  (dx,weights) = pair
73	  (dx,(weight_gradient,(bias_gradient,weights)))
74	
75	# The height and width of [N,C,H,W] or [C,H,W]; zero for another rank.
76	def image(shape: List<&2,U32>) -> U32 & U32:
77	  match shape:
78	    case [images,channels,height,width]: (height,width)
79	    case [channels,height,width]: (height,width)
80	    case shape: (0,0)
81	
82	def sized(+stride: U32,+padding: U32,weight_gradient: Tensor(),bias_gradient: Tensor(),dz: Tensor(),weights: Tensor(),size: U32 & U32) -> Tensor() & (Tensor() & (Tensor() & Tensor())):
83	  (height,width) = size
84	  finished(weight_gradient,bias_gradient,Gradient.input(dz,weights,stride,padding,height,width))
85	
86	def propagated(+stride: U32,+padding: U32,result: Parameters()) -> Tensor() & (Tensor() & (Tensor() & Tensor())):
87	  (weight_gradient,(bias_gradient,(dz,(weights,shape)))) = result
88	  sized(stride,padding,weight_gradient,bias_gradient,dz,weights,image(shape))
89	
90	# dy, the gradient of the loss with respect to the layer output:
91	# (the gradient of the layer input, (weight gradient, (bias gradient, weights))).
92	def backward(dy: Tensor(),saved: Saved(),weights: Tensor(),+settings: Settings) -> Tensor() & (Tensor() & (Tensor() & Tensor())):
93	  Settings{+stride,+padding,activation} = settings
94	  propagated(stride,padding,parameters(dy,saved,weights,settings))
95	
96	def without_input(result: Parameters()) -> Tensor() & (Tensor() & Tensor()):
97	  (weight_gradient,(bias_gradient,(dz,(weights,shape)))) = result
98	  (weight_gradient,(bias_gradient,weights))
99	
100	# backward for a layer whose input needs no gradient:
101	# (weight gradient, (bias gradient, weights)).
102	def backward_parameters(dy: Tensor(),saved: Saved(),weights: Tensor(),settings: Settings) -> Tensor() & (Tensor() & Tensor()):
103	  without_input(parameters(dy,saved,weights,settings))