tensor_band_view.bend
stelliferous@0.0.3.0 · 230 lines · 0 definitions
1# A view of the storage of a geometry, seen from its bands. A view that keeps
2# the images and the rows of the geometry in place has axes before the rows,
3# repeats of each row and axes after the rows, each moving by whole channels
4# and columns; every band reads its cells through the same axes over its own
5# rows.
6import Base
7import stelliferous@0.0.3.0/tensor_view.bend as Views
8import stelliferous@0.0.3.0/tensor_shape.bend as Shapes
9import stelliferous@0.0.3.0/tensor_band.bend as Bands
10
11# The axes before the rows, the repeats of each row, the axes after the rows,
12# and the offset of the view, as the view gives them.
13type Banded is Data:
14 Banded{before: List<&2,Views.Axis>,repeats: List<&2,U32>,after: List<&2,Views.Axis>,offset: U32}
15
16# ---- The view of a band ----
17
18# The channel of plane words a word of the storage of a geometry lies in, or
19# its column within that channel.
20def component(+plane: U32,column: Bool,+word: U32) -> U32:
21 match column:
22 case False{}: (word / plane : U32)
23 case True{}: (word % plane : U32)
24
25# A word of the storage of the geometry, a channel of plane words and a column
26# within it, as the same channel and column in a band of band words per
27# channel.
28def localized(+plane: U32,+band: U32,+word: U32) -> U32:
29 (component(plane,False{},word) * band + component(plane,True{},word) : U32)
30
31def localized_axes(+plane: U32,+band: U32,axes: List<&2,Views.Axis>) -> List<&2,Views.Axis>:
32 match axes:
33 case []: []
34 case Views.Axis{extent,stride} <> rest: Views.Axis{extent,localized(plane,band,stride)} <> localized_axes(plane,band,rest)
35
36def repeated(repeats: List<&2,U32>) -> List<&2,Views.Axis>:
37 match repeats:
38 case []: []
39 case extent <> rest: Views.Axis{extent,0} <> repeated(rest)
40
41# The rows axis, which a view of a geometry of one row per image leaves out,
42# and the repeats of each row.
43def rows_axes(rowed: Bool,+rows: U32,+width: U32,repeats: List<&2,U32>) -> List<&2,Views.Axis>:
44 match rowed:
45 case True{}: Views.Axis{rows,width} <> repeated(repeats)
46 case False{}: repeated(repeats)
47
48# The axes of the view over rows rows of each channel. A band of a geometry
49# of one row per image has at most one row.
50def axes(+banded: Banded,+geometry: Bands.Geometry,+rows: U32) -> List<&2,Views.Axis>:
51 Banded{before,repeats,after,offset} = banded
52 Bands.Geometry{batch,channels,+height,+width} = geometry
53 +plane = (height * width : U32)
54 +band = (rows * width : U32)
55 List.append(&2,Views.Axis,localized_axes(plane,band,before),List.append(&2,Views.Axis,rows_axes(True{},rows,width,repeats),localized_axes(plane,band,after)))
56
57# The view of the local storage of a band of rows rows.
58def local(+banded: Banded,+geometry: Bands.Geometry,+rows: U32) -> Views.View:
59 Banded{before,repeats,after,+offset} = banded
60 Bands.Geometry{batch,channels,+height,+width} = geometry
61 Views.View{axes(banded,geometry,rows),localized((height * width : U32),(rows * width : U32),offset)}
62
63# The extents of the view over rows rows.
64def extents(+banded: Banded,+rows: U32) -> List<&2,U32>:
65 Banded{before,repeats,after,offset} = banded
66 List.append(&2,U32,Views.axis_extents(before),List.append(&2,U32,rows <> repeats,Views.axis_extents(after)))
67
68def known(count: Maybe<&2,U32>) -> U32:
69 match count:
70 case Some{count}: count
71 case None{}: 0
72
73def product_or_zero(extents: List<&2,U32>) -> U32:
74 known(Shapes.product(extents,Some{1}))
75
76# The rows of the output for each row of the input.
77def scale(+banded: Banded) -> U32:
78 Banded{before,repeats,after,offset} = banded
79 product_or_zero(repeats)
80
81# The geometry of the output: the same images, the axes before the rows as
82# channels, each row repeated, the axes after the rows as columns.
83def output(+banded: Banded,+geometry: Bands.Geometry) -> Bands.Geometry:
84 Banded{before,repeats,after,offset} = banded
85 Bands.Geometry{batch,channels,height,width} = geometry
86 Bands.Geometry{batch,product_or_zero(Views.axis_extents(before)),(height * product_or_zero(repeats) : U32),product_or_zero(Views.axis_extents(after))}
87
88# ---- Directions ----
89
90def count_when(present: Bool) -> Nat:
91 match present:
92 case True{}: 1n
93 case False{}: 0n
94
95def local_when(valid: Bool,flags: List<&2,Bool>) -> Maybe<&2,List<&2,Bool>>:
96 match valid:
97 case True{}: Some{flags}
98 case False{}: None{}
99
100# The flags of the axes of a local view: those of the axes before the rows,
101# the rows axis and its repeats forward, those of the axes after the rows.
102def local_directions(+before: List<&2,Views.Axis>,+repeats: List<&2,U32>,+after: List<&2,Views.Axis>,+before_flags: List<&2,Bool>,+after_flags: List<&2,Bool>) -> List<&2,Bool>:
103 List.append(&2,Bool,Views.fitted_flags(List.length(&2,Views.Axis,before),before_flags),
104 List.append(&2,Bool,Views.fitted_flags(1n+List.length(&2,U32,repeats),[]),Views.fitted_flags(List.length(&2,Views.Axis,after),after_flags)))
105
106# The flags of the local view of a band, when the view reads its images axis,
107# its rows axis and the repeats of each row forward. An axis without a flag is
108# read forward.
109def local_flags(+banded: Banded,+geometry: Bands.Geometry,+flags: List<&2,Bool>) -> Maybe<&2,List<&2,Bool>>:
110 Banded{+before,+repeats,+after,offset} = banded
111 Bands.Geometry{+batch,channels,+height,width} = geometry
112 +images = count_when(U32.is_lt(1,batch))
113 +rows = Nat.add(count_when(U32.is_lt(1,height)),List.length(&2,U32,repeats))
114 +later = Views.drop_flags(Nat.add(images,List.length(&2,Views.Axis,before)),flags)
115 local_when(Views.forward(images,flags) && Views.forward(rows,later),local_directions(before,repeats,after,Views.drop_flags(images,flags),Views.drop_flags(rows,later)))
116
117# ---- Classification ----
118
119# The axes after the images axis, which a view of more than one image starts
120# with.
121def images_kept(kept: Bool,rest: List<&2,Views.Axis>) -> Maybe<&2,List<&2,Views.Axis>>:
122 match kept:
123 case True{}: Some{rest}
124 case False{}: None{}
125
126def after_images(many: Bool,+batch: U32,+image: U32,axes: List<&2,Views.Axis>) -> Maybe<&2,List<&2,Views.Axis>>:
127 match many axes:
128 case False{} axes: Some{axes}
129 case True{} Nil{}: None{}
130 case True{} Views.Axis{+extent,+stride} <> rest: images_kept(U32.is_eq(extent,batch) && U32.is_eq(stride,image),rest)
131
132# The axes before the rows axis and those after it.
133type Split is Data:
134 Split{before: List<&2,Views.Axis>,rest: List<&2,Views.Axis>}
135
136def prepend_before(axis: Views.Axis,split: Maybe<&2,Split>) -> Maybe<&2,Split>:
137 match split:
138 case Some{Split{before,rest}}: Some{Split{axis <> before,rest}}
139 case None{}: None{}
140
141def split_from(rest: List<&2,Views.Axis>,found: Bool,+height: U32,+width: U32,+extent: U32,+stride: U32) -> Maybe<&2,Split>:
142 match rest found:
143 case rest True{}: Some{Split{[],rest}}
144 case Nil{} False{}: None{}
145 case Views.Axis{+next_extent,+next_stride} <> tail False{}:
146 prepend_before(Views.Axis{extent,stride},split_from(tail,U32.is_eq(next_extent,height) && U32.is_eq(next_stride,width),height,width,next_extent,next_stride))
147
148def split_rows(+height: U32,+width: U32,axes: List<&2,Views.Axis>) -> Maybe<&2,Split>:
149 match axes:
150 case []: None{}
151 case Views.Axis{+extent,+stride} <> rest: split_from(rest,U32.is_eq(extent,height) && U32.is_eq(stride,width),height,width,extent,stride)
152
153# The leading axes of stride zero repeat each row.
154def prepend_repeat(extent: U32,split: List<&2,U32> & List<&2,Views.Axis>) -> List<&2,U32> & List<&2,Views.Axis>:
155 (repeats,rest) = split
156 (extent <> repeats,rest)
157
158def repeats_from(rest: List<&2,Views.Axis>,zero: Bool,extent: U32,stride: U32) -> List<&2,U32> & List<&2,Views.Axis>:
159 match rest zero:
160 case rest False{}: ([],Views.Axis{extent,stride} <> rest)
161 case Nil{} True{}: ([extent],[])
162 case Views.Axis{next_extent,+next_stride} <> tail True{}: prepend_repeat(extent,repeats_from(tail,U32.is_zero(next_stride),next_extent,next_stride))
163
164def split_repeats(axes: List<&2,Views.Axis>) -> List<&2,U32> & List<&2,Views.Axis>:
165 match axes:
166 case []: ([],[])
167 case Views.Axis{extent,+stride} <> rest: repeats_from(rest,U32.is_zero(stride),extent,stride)
168
169# The largest channel, or column, the axes reach from a start.
170def reach(+plane: U32,+column: Bool,axes: List<&2,Views.Axis>,start: Maybe<&2,U32>) -> Maybe<&2,U32>:
171 match axes start:
172 case [] start: start
173 case Views.Axis{+extent,stride} <> rest Some{+start}:
174 reach(plane,column,rest,Shapes.checked_add_to(Shapes.checked_multiply((extent - 1 : U32),component(plane,column,stride)),start))
175 case axes None{}: None{}
176
177def below(reach: Maybe<&2,U32>,+limit: U32) -> Bool:
178 match reach:
179 case Some{reach}: U32.is_lt(reach,limit)
180 case None{}: False{}
181
182def counted(extents: List<&2,U32>) -> Bool:
183 Views.has_count(Shapes.product(extents,Some{1}))
184
185# A classification is kept when the geometry has images and rows, every axis
186# stays inside the channels and the columns of a band and every group of
187# extents has a count.
188def kept_when(valid: Bool,banded: Banded) -> Maybe<&2,Banded>:
189 match valid:
190 case True{}: Some{banded}
191 case False{}: None{}
192
193def fits(+banded: Banded,+geometry: Bands.Geometry) -> Bool:
194 Banded{+before,+repeats,+after,+offset} = banded
195 Bands.Geometry{+batch,+channels,+height,+width} = geometry
196 +plane = (height * width : U32)
197 U32.is_lt(0,batch) && U32.is_lt(0,height) && below(reach(plane,False{},after,reach(plane,False{},before,Some{component(plane,False{},offset)})),channels) &&
198 below(reach(plane,True{},after,reach(plane,True{},before,Some{component(plane,True{},offset)})),width) &&
199 counted(Views.axis_extents(before)) && counted(Views.axis_extents(after)) && counted(repeats) && Views.has_count(Shapes.checked_multiply(height,product_or_zero(repeats)))
200
201def kept(+banded: Banded,+geometry: Bands.Geometry) -> Maybe<&2,Banded>:
202 kept_when(fits(banded,geometry),banded)
203
204def split_kept(+offset: U32,+geometry: Bands.Geometry,before: List<&2,Views.Axis>,split: List<&2,U32> & List<&2,Views.Axis>) -> Maybe<&2,Banded>:
205 (repeats,after) = split
206 kept(Banded{before,repeats,after,offset},geometry)
207
208def rowed_split(+offset: U32,+geometry: Bands.Geometry,split: Maybe<&2,Split>) -> Maybe<&2,Banded>:
209 match split:
210 case Some{Split{before,rest}}: split_kept(offset,geometry,before,split_repeats(rest))
211 case None{}: None{}
212
213def split_geometry(+geometry: Bands.Geometry,axes: List<&2,Views.Axis>) -> Maybe<&2,Split>:
214 Bands.Geometry{batch,channels,+height,+width} = geometry
215 split_rows(height,width,axes)
216
217def classify_rows(rowed: Bool,axes: Maybe<&2,List<&2,Views.Axis>>,+offset: U32,+geometry: Bands.Geometry) -> Maybe<&2,Banded>:
218 match rowed axes:
219 case True{} Some{axes}: rowed_split(offset,geometry,split_geometry(geometry,axes))
220 case False{} Some{axes}: kept(Banded{[],[],axes,offset},geometry)
221 case rowed None{}: None{}
222
223# The banded form of a view, given by its axes and offset, of the storage of a
224# geometry, when the view starts with the images axis of a batch of more than
225# one image, has the rows axis of a geometry of more than one row, and stays
226# inside the channels and columns. A geometry of one image of one row, as a
227# whole tensor has, takes every view whose addresses fit it.
228def classify(axes: List<&2,Views.Axis>,+offset: U32,+geometry: Bands.Geometry) -> Maybe<&2,Banded>:
229 Bands.Geometry{+batch,+channels,+height,+width} = geometry
230 classify_rows(U32.is_lt(1,height),after_images(U32.is_lt(1,batch),batch,(channels * (height * width : U32) : U32),axes),offset,geometry)