1	# The everyday F32 tensor module: import this one module and call F.add,
2	# F.concat, F.check, ... without element-type arguments. Every function is a
3	# direct specialization of the generic tensor modules; Base determines operand
4	# order, rounding, NaN and signed-zero behavior.
5	import Base
6	import ./storage_buffer.bend as Buffers
7	import ./math_activation.bend as Activations
8	import ./math_fp32.bend as Arithmetic
9	import ./math_compare.bend as Compare
10	import ./tensor.bend as Tensors
11	import ./tensor_operations.bend as Operations
12	import ./tensor_lanes.bend as Lanes
13	import ./tensor_select.bend as Select
14	import ./tensor_layout.bend as Layout
15	import ./tensor_window.bend as Window
16	import ./tensor_convolution.bend as TensorConvolution
17	import ./tensor_file.bend as Files
18	import ./tensor_model.bend as Models
19	
20	def Tensor() -> Type:
21	  Tensors.Tensor<F32>
22	
23	def Problem() -> Data:
24	  Tensors.Problem
25	
26	# Storage exchanged with foreign I/O; see from_storage and into_storage.
27	def Storage() -> Type:
28	  Buffers.Buffer<F32>
29	
30	# Construction
31	
32	def full(+shape: List<&2,U32>,+value: F32) -> Tensor():
33	  Tensors.full(~F32,shape,value)
34	
35	def zeros(+shape: List<&2,U32>) -> Tensor():
36	  Tensors.full(~F32,shape,0.0)
37	
38	def ones(+shape: List<&2,U32>) -> Tensor():
39	  full(shape,1.0)
40	
41	def from_list(+shape: List<&2,U32>,+values: List<&2,F32>) -> Tensor():
42	  Tensors.from_list(~F32,shape,values,0.0)
43	
44	def from_storage(+shape: List<&2,U32>,storage: Storage()) -> Tensor():
45	  Tensors.from_storage(~F32,shape,storage,0.0)
46	
47	# Raw little-endian FP32 files: the values of the shape in row-major order.
48	# A file that cannot be read, or does not match the shape, gives a failed tensor.
49	def load_raw(+path: String,+shape: List<&2,U32>) -> IO(Tensor()):
50	  Files.load(path,shape)
51	
52	def save_raw(+path: String,tensor: Tensor()) -> IO(Tensor()):
53	  Files.save(path,tensor)
54	
55	# NumPy .npy files of FP32 values; the file carries the shape.
56	def load_npy(+path: String) -> IO(Tensor()):
57	  Files.load_npy(path)
58	
59	def save_npy(+path: String,tensor: Tensor()) -> IO(Tensor()):
60	  Files.save_npy(path,tensor)
61	
62	# Models: tensors by name, loaded once. In `do Model<T>:` a step takes a tensor
63	# from the weights with take, or puts a copy there with keep for a later take.
64	def Weights() -> Type:
65	  Models.Store
66	
67	def Model(A: Type) -> Type:
68	  Models.Model(A)
69	
70	def Model.pure(-A: Type,value: A) -> Model(A):
71	  Models.Model.pure(A,value)
72	
73	def Model.bind(-A: Type,-B: Type,step: Model(A),next: A -> Model(B)) -> Model(B):
74	  Models.Model.bind(A,B,step,next)
75	
76	def take(+name: String) -> Model(Tensor()):
77	  Models.take(name)
78	
79	def keep(+name: String,tensor: Tensor()) -> Model(Tensor()):
80	  Models.keep(name,tensor)
81	
82	# folder/names.txt lists one name per line; folder/<name>.npy holds its tensor.
83	def load_weights(+folder: String) -> IO(Weights()):
84	  Models.load(folder)
85	
86	def copy_weights(weights: Weights()) -> Weights() & Weights():
87	  Models.copy(weights)
88	
89	def run(-A: Type,model: Model(A),weights: Weights()) -> A:
90	  Models.result(A,model,weights)
91	
92	# The result with what remains of the weights, kept tensors included.
93	def run_keeping(-A: Type,model: Model(A),weights: Weights()) -> A & Weights():
94	  Models.run(A,model,weights)
95	
96	def save_weights(+folder: String,weights: Weights()) -> IO(Weights()):
97	  Models.save(folder,weights)
98	
99	# Execution context: operations may use count workers, like a tensor's device.
100	# Read the runtime's pool once with IO.thread_count(); results keep the count.
101	def with_workers(tensor: Tensor(),+count: U32) -> Tensor():
102	  Tensors.with_workers(~F32,tensor,count)
103	
104	def workers(tensor: Tensor()) -> Tensor() & U32:
105	  Tensors.workers(~F32,tensor)
106	
107	# Observation and exits
108	
109	def shape(tensor: Tensor()) -> Tensor() & List<&2,U32>:
110	  Tensors.shape(~F32,tensor)
111	
112	def check(tensor: Tensor()) -> Result<&1,&1,Tensor() & Problem(),Tensor()>:
113	  Tensors.check(~F32,tensor)
114	
115	def describe(problem: Problem()) -> String:
116	  Tensors.describe(problem)
117	
118	def expect(tensor: Tensor()) -> IO(Tensor()):
119	  Tensors.expect(~F32,tensor)
120	
121	def to_list(tensor: Tensor()) -> Result<&2,&2,Problem(),List<&2,F32>>:
122	  Tensors.to_list(~F32,tensor)
123	
124	def item(tensor: Tensor()) -> Result<&2,&2,Problem(),F32>:
125	  Tensors.item(~F32,tensor)
126	
127	# Reads with a fallback for a failed tensor, for callers that do not inspect
128	# the problem.
129	def to_list_or(tensor: Tensor(),fallback: List<&2,F32>) -> List<&2,F32>:
130	  Tensors.value_or(~List<&2,F32>,Tensors.to_list(~F32,tensor),fallback)
131	
132	def item_or(tensor: Tensor(),fallback: F32) -> F32:
133	  Tensors.value_or(~F32,Tensors.item(~F32,tensor),fallback)
134	
135	def into_storage(tensor: Tensor()) -> Result<&2,&1,Problem(),Storage()>:
136	  Tensors.into_storage(~F32,tensor)
137	
138	def expect_storage(tensor: Tensor()) -> IO(Storage()):
139	  Tensors.expect_storage(~F32,tensor)
140	
141	# Operations consume their tensor arguments and return one result. Keep a
142	# tensor by copying it first; copy uses then, and shape uses then_shape.
143	
144	def then(pair: Tensor() & Tensor(),body: Tensor() -> Tensor() -> Tensor()) -> Tensor():
145	  Tensors.then(~Tensor(),~Tensor(),~Tensor(),pair,body)
146	
147	def then_shape(pair: Tensor() & List<&2,U32>,body: Tensor() -> List<&2,U32> -> Tensor()) -> Tensor():
148	  Tensors.then(~Tensor(),~List<&2,U32>,~Tensor(),pair,body)
149	
150	def copy(tensor: Tensor()) -> Tensor() & Tensor():
151	  Layout.copy(~F32,tensor)
152	
153	# Binary elementwise operations reuse the left tensor's storage. The number
154	# after the name is the operation's scalar work for the band schedule.
155	def add(left: Tensor(),right: Tensor()) -> Tensor():
156	  Tensors.combine(~F32,~F32.add,"add",1.0,left,right)
157	
158	def sub(left: Tensor(),right: Tensor()) -> Tensor():
159	  Tensors.combine(~F32,~F32.sub,"sub",1.0,left,right)
160	
161	def mul(left: Tensor(),right: Tensor()) -> Tensor():
162	  Tensors.combine(~F32,~F32.mul,"mul",1.0,left,right)
163	
164	# Base division: a zero divisor gives an infinity or NaN, never a failure.
165	def div(left: Tensor(),right: Tensor()) -> Tensor():
166	  Tensors.combine(~F32,~F32.div,"div",1.0,left,right)
167	
168	def maximum(left: Tensor(),right: Tensor()) -> Tensor():
169	  Tensors.combine(~F32,~F32.max,"maximum",2.0,left,right)
170	
171	def minimum(left: Tensor(),right: Tensor()) -> Tensor():
172	  Tensors.combine(~F32,~F32.min,"minimum",2.0,left,right)
173	
174	def add_scalar(tensor: Tensor(),+value: F32) -> Tensor():
175	  Tensors.combine_scalar(~F32,~F32.add,"add_scalar",1.0,tensor,value)
176	
177	def sub_scalar(tensor: Tensor(),+value: F32) -> Tensor():
178	  Tensors.combine_scalar(~F32,~F32.sub,"sub_scalar",1.0,tensor,value)
179	
180	def mul_scalar(tensor: Tensor(),+value: F32) -> Tensor():
181	  Tensors.combine_scalar(~F32,~F32.mul,"mul_scalar",1.0,tensor,value)
182	
183	def div_scalar(tensor: Tensor(),+value: F32) -> Tensor():
184	  Tensors.combine_scalar(~F32,~F32.div,"div_scalar",1.0,tensor,value)
185	
186	def less(left: Tensor(),right: Tensor()) -> Tensor():
187	  Tensors.combine_with(~F32,~Lanes.combine(~Compare.less_elements),"less",2.625,left,right)
188	
189	def less_equal(left: Tensor(),right: Tensor()) -> Tensor():
190	  Tensors.combine_with(~F32,~Lanes.combine(~Compare.less_equal_elements),"less_equal",2.625,left,right)
191	
192	def greater(left: Tensor(),right: Tensor()) -> Tensor():
193	  Tensors.combine_with(~F32,~Lanes.combine(~Compare.greater_elements),"greater",2.625,left,right)
194	
195	def greater_equal(left: Tensor(),right: Tensor()) -> Tensor():
196	  Tensors.combine_with(~F32,~Lanes.combine(~Compare.greater_equal_elements),"greater_equal",2.625,left,right)
197	
198	def equal(left: Tensor(),right: Tensor()) -> Tensor():
199	  Tensors.combine_with(~F32,~Lanes.combine(~Compare.equal_elements),"equal",1.375,left,right)
200	
201	def not_equal(left: Tensor(),right: Tensor()) -> Tensor():
202	  Tensors.combine_with(~F32,~Lanes.combine(~Compare.not_equal_elements),"not_equal",1.375,left,right)
203	
204	# Scalar comparisons share the eight-lane map and preserve the tensor layout.
205	# Costs count AVX2 arithmetic, comparison, selection and value-bit instructions
206	# in one eight-element round, divided by eight; moves and addressing are excluded.
207	def less_scalar(tensor: Tensor(),+value: F32) -> Tensor():
208	  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.less),value,"less_scalar",2.75,tensor)
209	
210	def less_equal_scalar(tensor: Tensor(),+value: F32) -> Tensor():
211	  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.less_equal),value,"less_equal_scalar",2.625,tensor)
212	
213	def greater_scalar(tensor: Tensor(),+value: F32) -> Tensor():
214	  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.greater),value,"greater_scalar",2.625,tensor)
215	
216	def greater_equal_scalar(tensor: Tensor(),+value: F32) -> Tensor():
217	  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.greater_equal),value,"greater_equal_scalar",2.75,tensor)
218	
219	def equal_scalar(tensor: Tensor(),+value: F32) -> Tensor():
220	  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.equal),value,"equal_scalar",1.5,tensor)
221	
222	def not_equal_scalar(tensor: Tensor(),+value: F32) -> Tensor():
223	  Tensors.apply_context(~F32,~F32,~Operations.lanes(~F32,~Compare.not_equal),value,"not_equal_scalar",1.5,tensor)
224	
225	def select(mask: Tensor(),yes: Tensor(),no: Tensor()) -> Tensor():
226	  Select.select(mask,yes,no)
227	
228	def clamp(tensor: Tensor(),+low: F32,+high: F32) -> Tensor():
229	  Tensors.apply_context(~F32,~Compare.Bounds,~Operations.lanes(~Compare.Bounds,~Compare.clamp),Compare.Bounds{low,high},"clamp",0.5,tensor)
230	
231	# Activations: Act.silu() and the other values of lib/math_activation.bend
232	# carry their own planner cost. activate applies one to a tensor; relu, silu,
233	# ... below are its shorthands. Custom carries only a cost; its function
234	# reaches conv2d_custom or apply, and activate fails on it.
235	def activate(tensor: Tensor(),+activation: Activations.Activation) -> Tensor():
236	  match activation:
237	    case Activations.Identity{}: tensor
238	    case Activations.ReLU{}: Tensors.apply(~F32,~Operations.mapped(~F32,~Activations.relu_value),"relu",Activations.operations(Activations.ReLU{}),tensor)
239	    case Activations.ReLU6{}: Tensors.apply(~F32,~Operations.mapped(~F32,~Activations.relu6_value),"relu6",Activations.operations(Activations.ReLU6{}),tensor)
240	    case Activations.LeakyReLU{+slope}: Tensors.combine_scalar(~F32,~Activations.leaky_relu_value,"leaky_relu",Activations.operations(Activations.LeakyReLU{slope}),tensor,slope)
241	    case Activations.Sigmoid{}: Tensors.apply(~F32,~Operations.mapped_lanes(~Activations.sigmoid_value),"sigmoid",Activations.operations(Activations.Sigmoid{}),tensor)
242	    case Activations.SiLU{}: Tensors.apply(~F32,~Operations.mapped_lanes(~Activations.silu_value),"silu",Activations.operations(Activations.SiLU{}),tensor)
243	    case Activations.Tanh{}: Tensors.apply(~F32,~Operations.mapped_lanes(~Activations.tanh_value),"tanh",Activations.operations(Activations.Tanh{}),tensor)
244	    case Activations.HardSigmoid{}: Tensors.apply(~F32,~Operations.mapped(~F32,~Activations.hard_sigmoid_value),"hard_sigmoid",Activations.operations(Activations.HardSigmoid{}),tensor)
245	    case Activations.HardSwish{}: Tensors.apply(~F32,~Operations.mapped(~F32,~Activations.hard_swish_value),"hard_swish",Activations.operations(Activations.HardSwish{}),tensor)
246	    case Activations.GELUTanh{}: Tensors.apply(~F32,~Operations.mapped_lanes(~Activations.gelu_tanh_value),"gelu_tanh",Activations.operations(Activations.GELUTanh{}),tensor)
247	    case Activations.Custom{count}: Tensors.rejected(~F32,"activate",tensor)
248	
249	def relu(tensor: Tensor()) -> Tensor():
250	  activate(tensor,Activations.relu())
251	
252	def relu6(tensor: Tensor()) -> Tensor():
253	  activate(tensor,Activations.relu6())
254	
255	def leaky_relu(tensor: Tensor(),slope: F32) -> Tensor():
256	  activate(tensor,Activations.leaky_relu(slope))
257	
258	def sigmoid(tensor: Tensor()) -> Tensor():
259	  activate(tensor,Activations.sigmoid())
260	
261	def silu(tensor: Tensor()) -> Tensor():
262	  activate(tensor,Activations.silu())
263	
264	def tanh(tensor: Tensor()) -> Tensor():
265	  activate(tensor,Activations.tanh())
266	
267	def hard_sigmoid(tensor: Tensor()) -> Tensor():
268	  activate(tensor,Activations.hard_sigmoid())
269	
270	def hard_swish(tensor: Tensor()) -> Tensor():
271	  activate(tensor,Activations.hard_swish())
272	
273	# PyTorch's gelu(approximate='tanh').
274	def gelu_tanh(tensor: Tensor()) -> Tensor():
275	  activate(tensor,Activations.gelu_tanh())
276	
277	# Ordered convolution over CHW or NCHW; the activation runs on each finished
278	# leaf while it is in cache.
279	def conv2d(input: Tensor(),weights: Tensor(),bias: Tensor(),stride: U32,padding: U32,activation: Activations.Activation) -> Tensor():
280	  TensorConvolution.conv2d(~Activations.identity_value,activation,input,weights,bias,stride,padding)
281	
282	# conv2d with a custom activation of the given scalar operations per element.
283	def conv2d_custom(~activation: F32 -> F32,operations: F32,input: Tensor(),weights: Tensor(),bias: Tensor(),stride: U32,padding: U32) -> Tensor():
284	  TensorConvolution.conv2d(~activation,Activations.Custom{operations},input,weights,bias,stride,padding)
285	
286	# Custom elementwise functions are passed as direct templates, never records,
287	# with their scalar operations per element.
288	def apply(~operation: F32 -> F32,operations: F32,tensor: Tensor()) -> Tensor():
289	  Tensors.apply(~F32,~Operations.mapped(~F32,~operation),"apply",operations,tensor)
290	
291	# The function also receives the element's row-major index.
292	def apply_indexed(~operation: U32 -> F32 -> F32,operations: F32,tensor: Tensor()) -> Tensor():
293	  Tensors.apply_indexed(~F32,~operation,operations,tensor)
294	
295	# start, start + 1, ... as a rank-1 tensor of count elements.
296	def arange(+start: F32,+count: U32) -> Tensor():
297	  add_scalar(Tensors.apply_indexed(~F32,~(index => value => U32.to_f32(index)),1.0,zeros([count])),start)
298	
299	# Reductions return a rank-0 tensor, so failures propagate like any result.
300	
301	def sum(tensor: Tensor()) -> Tensor():
302	  Tensors.reduce_all(~F32,~F32.add,"sum",tensor,0.0)
303	
304	def max(tensor: Tensor()) -> Tensor():
305	  Tensors.reduce_all(~F32,~F32.max,"max",tensor,F32.div(F32.neg(1.0),0.0))
306	
307	def min(tensor: Tensor()) -> Tensor():
308	  Tensors.reduce_all(~F32,~F32.min,"min",tensor,F32.div(1.0,0.0))
309	
310	# sum divided by the number of elements; an empty tensor gives NaN.
311	def mean(tensor: Tensor()) -> Tensor():
312	  then_shape(shape(tensor),tensor => extents =>
313	    Tensors.combine_scalar(~F32,~F32.div,"mean",1.0,Tensors.reduce_all(~F32,~F32.add,"mean",tensor,0.0),
314	      U32.to_f32(Tensors.count_or_zero(extents))))
315	
316	# Layout: results are new dense tensors, except reshape, which relabels.
317	
318	def reshape(tensor: Tensor(),+shape: List<&2,U32>) -> Tensor():
319	  Layout.reshape(~F32,tensor,shape)
320	
321	# The result's axis i is the argument's axis order[i].
322	def transpose(tensor: Tensor(),order: List<&2,U32>) -> Tensor():
323	  Layout.transpose(~F32,tensor,order)
324	
325	# count elements from start along axis, every step-th one.
326	def slice(tensor: Tensor(),axis: U32,start: U32,count: U32,step: U32) -> Tensor():
327	  Layout.slice(~F32,tensor,axis,start,count,step)
328	
329	# target with part written at start along axis.
330	def set_slice(target: Tensor(),axis: U32,start: U32,part: Tensor()) -> Tensor():
331	  Layout.set_slice(~F32,target,axis,start,part)
332	
333	def Parts() -> Type:
334	  Tensors.Parts<F32>
335	
336	# Build the parts of concat: F.cons(a, F.cons(b, F.single(c))).
337	def single(tensor: Tensor()) -> Parts():
338	  Tensors.Single{tensor}
339	
340	def cons(tensor: Tensor(),rest: Parts()) -> Parts():
341	  Tensors.Cons{tensor,rest}
342	
343	def concat(parts: Parts(),+axis: U32) -> Tensor():
344	  Layout.concat(~F32,parts,axis)
345	
346	# Repeats each element of the last two axes scale times in each direction.
347	def upsample_nearest(tensor: Tensor(),+scale: U32) -> Tensor():
348	  Layout.upsample_nearest(~F32,tensor,scale)
349	
350	# Over the last two axes; positions outside the input count as -inf.
351	def max_pool2d_with(tensor: Tensor(),+kernel: U32,+stride: U32,+padding: U32) -> Tensor():
352	  Window.pool2d(~F32,~F32.max,"max_pool2d",F32.div(F32.neg(1.0),0.0),tensor,kernel,stride,padding)
353	
354	# Windows of kernel by kernel that do not overlap: stride kernel, padding 0.
355	def max_pool2d(tensor: Tensor(),+kernel: U32) -> Tensor():
356	  max_pool2d_with(tensor,kernel,kernel,0)
357	
358	# Broadcasting and axis operations.
359	
360	# The tensor repeated to shape under NumPy broadcasting; add, sub, mul, max
361	# and min broadcast their operands in the same way.
362	def broadcast_to(tensor: Tensor(),+shape: List<&2,U32>) -> Tensor():
363	  Layout.broadcast_to(~F32,tensor,shape)
364	
365	# Reductions along one axis remove that axis. Each result element takes the
366	# elements of the axis in increasing index, starting from 0, -inf or +inf.
367	def sum_axis(tensor: Tensor(),+axis: U32) -> Tensor():
368	  Layout.reduce_axis(~F32,~F32.add,"sum_axis",0.0,tensor,axis)
369	
370	def max_axis(tensor: Tensor(),+axis: U32) -> Tensor():
371	  Layout.reduce_axis(~F32,~F32.max,"max_axis",F32.div(F32.neg(1.0),0.0),tensor,axis)
372	
373	def min_axis(tensor: Tensor(),+axis: U32) -> Tensor():
374	  Layout.reduce_axis(~F32,~F32.min,"min_axis",F32.div(1.0,0.0),tensor,axis)
375	
376	# Elementwise functions share the eight-value traversal used by activations.
377	def exp(tensor: Tensor()) -> Tensor():
378	  Tensors.apply(~F32,~Operations.mapped_lanes(~Arithmetic.exp),"exp",4.5,tensor)
379	
380	def log(tensor: Tensor()) -> Tensor():
381	  Tensors.apply(~F32,~Operations.mapped_lanes(~Arithmetic.log),"log",8.25,tensor)
382	
383	# Keep a singleton reduction axis for broadcasting. Both reductions use the
384	# existing ordered fold; reduction gathers row bands to a dense owner.
385	def softmax(tensor: Tensor(),+axis: U32) -> Tensor():
386	  then_shape(shape(tensor),tensor => extents =>
387	    +unit = Layout.with_extent(extents,U32.to_nat(axis),1)
388	    then(copy(tensor),values => maximum_input =>
389	      then(copy(exp(sub(values,reshape(Layout.reduce_axis(~F32,~F32.max,"softmax",F32.div(F32.neg(1.0),0.0),maximum_input,axis),unit)))),
390	        values => sum_input => div(values,reshape(Layout.reduce_axis(~F32,~F32.add,"softmax",0.0,sum_input,axis),unit)))))
391	
392	# sum_axis divided by the extent of the axis.
393	def mean_axis(tensor: Tensor(),+axis: U32) -> Tensor():
394	  then_shape(shape(tensor),tensor => extents =>
395	    Tensors.combine_scalar(~F32,~F32.div,"mean_axis",1.0,Layout.reduce_axis(~F32,~F32.add,"mean_axis",0.0,tensor,axis),
396	      U32.to_f32(Layout.extent_at(extents,U32.to_nat(axis)))))
397	
398	# before[i] new elements in front of axis i and after[i] behind it, all equal
399	# to value.
400	def pad(tensor: Tensor(),+before: List<&2,U32>,+after: List<&2,U32>,+value: F32) -> Tensor():
401	  Layout.pad(~F32,Layout.Constant{value},tensor,before,after)
402	
403	# The same, each new element repeating the nearest element of the tensor.
404	def pad_edge(tensor: Tensor(),+before: List<&2,U32>,+after: List<&2,U32>) -> Tensor():
405	  Layout.pad(~F32,Layout.Edge{},tensor,before,after)