1	# Losses with their gradient, in a fixed FP32 order. The prediction becomes
2	# the gradient, and the target returns.
3	import Base
4	import ./tensor.bend as Tensors
5	import ./tensor_shape.bend as Shapes
6	import ./tensor_layout.bend as Layout
7	import ./tensor_operations.bend as Operations
8	import ./math_sum.bend as Sum
9	import ./tensor_f32.bend as F
10	import ./math_fp32.bend as Arithmetic
11	
12	def Tensor() -> Type:
13	  Tensors.Tensor<F32>
14	
15	type Reduction is Data:
16	  Mean{}
17	  Sum{}
18	
19	def mean() -> Reduction:
20	  Mean{}
21	
22	def sum() -> Reduction:
23	  Sum{}
24	
25	def square(value: F32) -> F32:
26	  +value = value
27	  (value * value : F32)
28	
29	def reduced(reduction: Reduction,+count: F32,total: Tensor()) -> Tensor():
30	  match reduction:
31	    case Mean{}: Tensors.combine_scalar(~F32,~F32.div,1.0,total,count)
32	    case Sum{}: total
33	
34	def mse_squares(+reduction: Reduction,+count: F32,target: Tensor(),pair: Tensor() & Tensor()) -> Tensor() & (Tensor() & Tensor()):
35	  (difference,squares) = pair
36	  (reduced(reduction,count,Layout.reduce_all(~F32,~Sum.Lanes,~Sum.step,~Sum.total,"mse_loss",6.0,Sum.zero(),
37	      Tensors.apply(~F32,~Operations.mapped(~F32,~square),1.0,squares))),
38	    (reduced(reduction,count,Tensors.combine_scalar(~F32,~F32.mul,1.0,difference,2.0)),target))
39	
40	def mse_counted(+reduction: Reduction,target: Tensor(),shaped: Tensor() & List<&2,U32>) -> Tensor() & (Tensor() & Tensor()):
41	  (difference,+shape) = shaped
42	  mse_squares(reduction,U32.to_f32(Tensors.count_or_zero(shape)),target,Layout.copy(~F32,difference))
43	
44	def mse_difference(same: Bool,prediction: Tensor(),kept: Tensor()) -> Tensor():
45	  match same:
46	    case True{}: Tensors.combine(~F32,~F32.sub,"mse_loss",1.0,prediction,kept)
47	    case False{}: Tensors.rejected(~F32,"mse_loss",prediction)
48	
49	def mse_kept(+reduction: Reduction,same: Bool,prediction: Tensor(),pair: Tensor() & Tensor()) -> Tensor() & (Tensor() & Tensor()):
50	  (target,kept) = pair
51	  mse_counted(reduction,target,Tensors.shape(~F32,mse_difference(same,prediction,kept)))
52	
53	def mse_shaped(+reduction: Reduction,shaped: Tensor() & List<&2,U32>,measured: Tensor() & List<&2,U32>) -> Tensor() & (Tensor() & Tensor()):
54	  (prediction,shape) = shaped
55	  (target,target_shape) = measured
56	  mse_kept(reduction,Shapes.same_extents(shape,target_shape),prediction,Layout.copy(~F32,target))
57	
58	# The loss of prediction against a target of the same shape, the gradient of
59	# the loss with respect to prediction, and the target. With d = prediction -
60	# target: Mean gives sum(d * d) / count and (d * 2) / count, Sum gives
61	# sum(d * d) and d * 2; the squares are added in row-major order from +0.
62	def mse(prediction: Tensor(),target: Tensor(),+reduction: Reduction) -> Tensor() & (Tensor() & Tensor()):
63	  mse_shaped(reduction,Tensors.shape(~F32,prediction),Tensors.shape(~F32,target))
64	
65	# ---- Cross entropy ----
66	
67	# The class numbers [N] as rows [N,K] of a target: 1 where the column equals
68	# the class number, 0 elsewhere.
69	def one_hot_rows(shape: List<&2,U32>,labels: Tensor(),+classes: U32) -> Tensor():
70	  match shape:
71	    case [images]: F.equal(F.reshape(labels,[images,1]),F.reshape(F.arange(0.0,classes),[1,classes]))
72	    case shape: Tensors.rejected(~F32,"one_hot",labels)
73	
74	def one_hot_measured(measured: Tensor() & List<&2,U32>,+classes: U32) -> Tensor():
75	  (labels,shape) = measured
76	  one_hot_rows(shape,labels,classes)
77	
78	def one_hot(labels: Tensor(),+classes: U32) -> Tensor():
79	  one_hot_measured(Tensors.shape(~F32,labels),classes)
80	
81	# The exponentials of a row added in the lanes of math_sum.bend, and the
82	# logarithm of their sum.
83	def exponential_step(lanes: Sum.Lanes,value: F32) -> Sum.Lanes:
84	  Sum.step(lanes,Arithmetic.exp(value))
85	
86	def logarithm_total(lanes: Sum.Lanes) -> F32:
87	  Arithmetic.log(Sum.total(lanes))
88	
89	def entropy_gradient(+images: U32,log_probabilities: Tensor() & Tensor(),targets: Tensor() & Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
90	  (for_loss,for_gradient) = log_probabilities
91	  (for_product,for_difference) = targets
92	  (F.div_scalar(F.sum(F.mul(for_product,for_loss)),F32.neg(U32.to_f32(images))),
93	    (F.div_scalar(F.sub(F.exp(for_gradient),for_difference),U32.to_f32(images)),target))
94	
95	def entropy_shifted(+images: U32,shifted: Tensor() & Tensor(),targets: Tensor() & Tensor()) -> Tensor() & (Tensor() & Tensor()):
96	  (values,for_sum) = shifted
97	  (target,target_kept) = targets
98	  entropy_gradient(images,
99	    F.copy(F.sub(values,F.reshape(Layout.reduce_axes(~F32,~Sum.Lanes,~exponential_step,~logarithm_total,"cross_entropy",10.5,Sum.zero(),for_sum,[1]),[images,1]))),
100	    F.copy(target),target_kept)
101	
102	def entropy_rows(+images: U32,logits: Tensor() & Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
103	  (values,for_maximum) = logits
104	  entropy_shifted(images,
105	    F.copy(F.sub(values,F.reshape(Layout.reduce_axes(~F32,~F32,~F32.max,~Tensors.same(~F32),"cross_entropy",1.0,F32.div(F32.neg(1.0),0.0),for_maximum,[1]),[images,1]))),
106	    F.copy(target))
107	
108	def entropy_failed(shape: List<&2,U32>,logits: Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
109	  (Tensors.failed(~F32,"cross_entropy",Tensors.InvalidShape{shape},[],0.0,1),(Tensors.rejected(~F32,"cross_entropy",logits),target))
110	
111	def entropy_ranked(shape: List<&2,U32>,logits: Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
112	  match shape:
113	    case [images,classes]: entropy_rows(images,F.copy(logits),target)
114	    case shape: entropy_failed(shape,logits,target)
115	
116	def entropy_shaped(same: Bool,logits: Tensor(),shape: List<&2,U32>,target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
117	  match same:
118	    case True{}: entropy_ranked(shape,logits,target)
119	    case False{}: entropy_failed(shape,logits,target)
120	
121	def entropy_measured(measured_logits: Tensor() & List<&2,U32>,measured_target: Tensor() & List<&2,U32>) -> Tensor() & (Tensor() & Tensor()):
122	  (logits,+shape) = measured_logits
123	  (target,target_shape) = measured_target
124	  entropy_shaped(Shapes.same_extents(shape,target_shape),logits,shape,target)
125	
126	# logits [N,K] and a target [N,K] whose rows are class probabilities, such as
127	# one_hot of the class numbers: the mean loss, its gradient with respect to
128	# logits, and the target. With m the maximum of a row and the log
129	# probabilities p = (logits - m) - log(sum of exp(logits - m) over the row),
130	# the sum over the row in the lanes of math_sum.bend: the loss is
131	# sum(target * p) / -N, the sum of math_sum.bend in row-major order, and the
132	# gradient (exp(p) - target) / N.
133	def cross_entropy(logits: Tensor(),target: Tensor()) -> Tensor() & (Tensor() & Tensor()):
134	  entropy_measured(Tensors.shape(~F32,logits),Tensors.shape(~F32,target))