kernel_shape.bend

stelliferous@0.0.2.0 · 47 lines · 7 definitions

1# Natural dimensions at the validation boundary, before any machine narrowing.
2# Convolution specifications and public tensor metadata use the same envelope.
3import Base
4import stelliferous@0.0.2.0/machine_limits.bend as Machine
5
6type Dimensions is Data:
7  Dimensions{input_channels: Nat,output_channels: Nat,height: Nat,width: Nat,kernel: Nat,stride: Nat,padding: Nat}
8
9def axis_size(+size: Nat,+kernel: Nat,stride: Nat,+padding: Nat) -> Nat:
10  1n+Nat.div(Nat.sub(Nat.add(size,Nat.mul(2n,padding)),kernel),stride)
11
12def valid(shape: Dimensions,input_length: Nat,weight_length: Nat,bias_length: Nat) -> Bool:
13  Dimensions{+ci,+co,+height,+width,+kernel,+stride,+padding} = shape
14  Nat.is_lt(0n,ci) && Nat.is_lt(0n,co) && Nat.is_lt(0n,height) && Nat.is_lt(0n,width) && Nat.is_lt(0n,kernel) && Nat.is_lt(0n,stride) &&
15  Nat.is_le(kernel,Nat.add(height,Nat.mul(2n,padding))) && Nat.is_le(kernel,Nat.add(width,Nat.mul(2n,padding))) &&
16  Nat.is_eq(input_length,Nat.mul(ci,Nat.mul(height,width))) &&
17  Nat.is_eq(weight_length,Nat.mul(co,Nat.mul(ci,Nat.mul(kernel,kernel)))) && Nat.is_eq(bias_length,co)
18
19def direct_bounds(+dimension: U32,+capacity: U32,+coordinate: U32,shape: Dimensions,input_length: Nat,weight_length: Nat) -> Bool:
20  Dimensions{+ci,+co,+height,+width,+kernel,+stride,+padding} = shape
21  +oh = axis_size(height,kernel,stride,padding)
22  +ow = axis_size(width,kernel,stride,padding)
23  Nat.is_le(ci,U32.to_nat(dimension)) && Nat.is_le(co,U32.to_nat(dimension)) &&
24  Nat.is_le(height,U32.to_nat(dimension)) && Nat.is_le(width,U32.to_nat(dimension)) &&
25  Nat.is_le(kernel,U32.to_nat(dimension)) && Nat.is_le(stride,U32.to_nat(dimension)) && Nat.is_le(padding,U32.to_nat(dimension)) &&
26  Nat.is_le(input_length,U32.to_nat(capacity)) && Nat.is_le(weight_length,U32.to_nat(capacity)) &&
27  Nat.is_le(Nat.mul(co,Nat.mul(oh,ow)),U32.to_nat(capacity)) &&
28  Nat.is_le(Nat.add(Nat.mul(oh,stride),kernel),U32.to_nat(coordinate)) &&
29  Nat.is_le(Nat.add(Nat.mul(ow,stride),kernel),U32.to_nat(coordinate))
30
31def packed_bounds(dimension: U32,+capacity: U32,coordinate: U32,+shape: Dimensions,input_length: Nat,weight_length: Nat) -> Bool:
32  Dimensions{+ci,+co,+height,+width,+kernel,+stride,+padding} = shape
33  +oh = axis_size(height,kernel,stride,padding)
34  +ow = axis_size(width,kernel,stride,padding)
35  +padded = Nat.mul(Nat.div(Nat.add(Nat.mul(oh,ow),7n),8n),8n)
36  direct_bounds(dimension,capacity,coordinate,shape,input_length,weight_length) &&
37  Nat.is_le(Nat.mul(padded,Nat.mul(ci,Nat.mul(kernel,kernel))),U32.to_nat(capacity)) &&
38  Nat.is_le(Nat.mul(padded,co),U32.to_nat(capacity))
39
40# Never duplicate the Limits owner: its Facts field depends on its values.
41def supports_direct(limits: Machine.Limits,shape: Dimensions,input_length: Nat,weight_length: Nat) -> Bool:
42  Machine.Limits{dimension,capacity,coordinate,half_word,facts} = limits
43  direct_bounds(dimension,capacity,coordinate,shape,input_length,weight_length)
44
45def supports_packed(limits: Machine.Limits,shape: Dimensions,input_length: Nat,weight_length: Nat) -> Bool:
46  Machine.Limits{dimension,capacity,coordinate,half_word,facts} = limits
47  packed_bounds(dimension,capacity,coordinate,shape,input_length,weight_length)