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