1	import Base
2	import ../shape/core.bend as Shape
3	import ./core.bend as Core
4	import ./transforms.bend as Transforms
5	import ./observation.bend as Observation
6	
7	type Request is Data:
8	  Permute{axes: List<&2, U32>}
9	  Transpose{a: U32, b: U32}
10	
11	def apply(+input: Shape.Shape, tensor: Core.Tensor<input>, request: Request) ->
12	  Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>:
13	  match request:
14	    case Permute{+axes}:
15	      Transforms.permute_copy(input, tensor, axes)
16	    case Transpose{+a, +b}:
17	      Transforms.transpose_copy(input, tensor, a, b)
18	
19	def status(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> U32:
20	  match result:
21	    case Fail{error}:
22	      Transforms.error_code(error)
23	    case Done{output}:
24	      0
25	
26	def failure_code_pair(-input: Shape.Shape, pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> U32:
27	  (source, result) = pair
28	  status(result)
29	
30	def failure_code(+input: Shape.Shape, +request: Request, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, U32>:
31	  match result:
32	    case Fail{error}:
33	      None{}
34	    case Done{tensor}:
35	      Some{failure_code_pair(input, apply(input, tensor, request))}
36	
37	def kept_value_pair(-input: Shape.Shape, +other: U32, pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, F32>:
38	  (source, result) = pair
39	  Core.read_value(input, Core.read_flat(input, source, other))
40	
41	def kept_value(+input: Shape.Shape, +request: Request, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, F32>:
42	  match result:
43	    case Fail{error}:
44	      None{}
45	    case Done{tensor}:
46	      kept_value_pair(input, other, apply(input, tensor, request))
47	
48	def kept_dimensions_pair(+input: Shape.Shape, pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, List<&2, U32>>:
49	  (source, result) = pair
50	  Observation.dimensions_pair(Core.shape(Core.pack(input, source)))
51	
52	def kept_dimensions(+input: Shape.Shape, +request: Request, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, List<&2, U32>>:
53	  match result:
54	    case Fail{error}:
55	      None{}
56	    case Done{tensor}:
57	      kept_dimensions_pair(input, apply(input, tensor, request))
58	
59	def kept_length_pair(-input: Shape.Shape, pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Nat:
60	  (source, result) = pair
61	  Core.logical_length(input, source)
62	
63	def kept_length(+input: Shape.Shape, +request: Request, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, Nat>:
64	  match result:
65	    case Fail{error}:
66	      None{}
67	    case Done{tensor}:
68	      Some{kept_length_pair(input, apply(input, tensor, request))}
69	
70	def kept_count_pair(+input: Shape.Shape, pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, U32>:
71	  (source, result) = pair
72	  Observation.count_pair(Core.count(Core.pack(input, source)))
73	
74	def kept_count(+input: Shape.Shape, +request: Request, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, U32>:
75	  match result:
76	    case Fail{error}:
77	      None{}
78	    case Done{tensor}:
79	      kept_count_pair(input, apply(input, tensor, request))
80	
81	def reshape_code(+input: Shape.Shape, +dimensions: List<&2, U32>, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, U32>:
82	  match result:
83	    case Fail{error}:
84	      None{}
85	    case Done{tensor}:
86	      Some{status(Transforms.reshape(input, tensor, dimensions))}
87	
88	def output_dimensions(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, List<&2, U32>>:
89	  match result:
90	    case Fail{error}:
91	      None{}
92	    case Done{output}:
93	      Some{Transforms.dynamic_dimensions(output)}
94	
95	def reshape_dimensions(+input: Shape.Shape, +dimensions: List<&2, U32>, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, List<&2, U32>>:
96	  match result:
97	    case Fail{error}:
98	      None{}
99	    case Done{tensor}:
100	      output_dimensions(Transforms.reshape(input, tensor, dimensions))
101	
102	def dynamic_length(dynamic: Core.Dynamic) -> Nat:
103	  Core.Packed{shape, tensor} = dynamic
104	  Core.logical_length(shape, tensor)
105	
106	def output_length(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, Nat>:
107	  match result:
108	    case Fail{error}:
109	      None{}
110	    case Done{output}:
111	      Some{dynamic_length(output)}
112	
113	def reshape_length(+input: Shape.Shape, +dimensions: List<&2, U32>, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, Nat>:
114	  match result:
115	    case Fail{error}:
116	      None{}
117	    case Done{tensor}:
118	      output_length(Transforms.reshape(input, tensor, dimensions))
119	
120	def output_count(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, U32>:
121	  match result:
122	    case Fail{error}:
123	      None{}
124	    case Done{output}:
125	      Observation.count_pair(Core.count(output))
126	
127	def reshape_count(+input: Shape.Shape, +dimensions: List<&2, U32>, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, U32>:
128	  match result:
129	    case Fail{error}:
130	      None{}
131	    case Done{tensor}:
132	      output_count(Transforms.reshape(input, tensor, dimensions))
133	
134	def output_value(+other: U32, result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, F32>:
135	  match result:
136	    case Fail{error}:
137	      None{}
138	    case Done{output}:
139	      Core.dynamic_read_value(Core.read_dynamic(output, other))
140	
141	def reshape_value(+input: Shape.Shape, +dimensions: List<&2, U32>, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, F32>:
142	  match result:
143	    case Fail{error}:
144	      None{}
145	    case Done{tensor}:
146	      output_value(other, Transforms.reshape(input, tensor, dimensions))
147	
148	def reshape_written_pair(+input: Shape.Shape, +dimensions: List<&2, U32>, +other: U32,
149	  pair: Core.Tensor<input> & Result<&2, &2, Core.Error, Unit>) -> Maybe<&2, F32>:
150	  (tensor, result) = pair
151	  output_value(other, Transforms.reshape(input, tensor, dimensions))
152	
153	def reshape_written(+input: Shape.Shape, +index: U32, +value: F32, +dimensions: List<&2, U32>, +other: U32,
154	  result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, F32>:
155	  match result:
156	    case Fail{error}:
157	      None{}
158	    case Done{tensor}:
159	      reshape_written_pair(input, dimensions, other, Core.write_flat(input, tensor, index, value))
160	
161	def applied_dimensions_pair(-input: Shape.Shape, pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, List<&2, U32>>:
162	  (source, result) = pair
163	  output_dimensions(result)
164	
165	def applied_dimensions(+input: Shape.Shape, +request: Request, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, List<&2, U32>>:
166	  match result:
167	    case Fail{error}:
168	      None{}
169	    case Done{tensor}:
170	      applied_dimensions_pair(input, apply(input, tensor, request))
171	
172	def applied_length_pair(-input: Shape.Shape, pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, Nat>:
173	  (source, result) = pair
174	  output_length(result)
175	
176	def applied_length(+input: Shape.Shape, +request: Request, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, Nat>:
177	  match result:
178	    case Fail{error}:
179	      None{}
180	    case Done{tensor}:
181	      applied_length_pair(input, apply(input, tensor, request))
182	
183	def applied_count_pair(-input: Shape.Shape, pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, U32>:
184	  (source, result) = pair
185	  output_count(result)
186	
187	def applied_count(+input: Shape.Shape, +request: Request, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, U32>:
188	  match result:
189	    case Fail{error}:
190	      None{}
191	    case Done{tensor}:
192	      applied_count_pair(input, apply(input, tensor, request))
193	
194	def applied_value_pair(-input: Shape.Shape, +other: U32, pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Maybe<&2, F32>:
195	  (source, result) = pair
196	  output_value(other, result)
197	
198	def applied_value(+input: Shape.Shape, +request: Request, +other: U32, result: Result<&2, &1, Core.Error, Core.Tensor<input>>) -> Maybe<&2, F32>:
199	  match result:
200	    case Fail{error}:
201	      None{}
202	    case Done{tensor}:
203	      applied_value_pair(input, other, apply(input, tensor, request))