(&self, context: Context, arguments_types: Vec<Type>)
| 49 | #[typetag::serde] |
| 50 | impl CustomOperationBody for Clip2K { |
| 51 | fn instantiate(&self, context: Context, arguments_types: Vec<Type>) -> Result<Graph> { |
| 52 | if arguments_types.len() != 1 { |
| 53 | return Err(runtime_error!("Invalid number of arguments for Clip")); |
| 54 | } |
| 55 | let input_type = arguments_types[0].clone(); |
| 56 | if !input_type.is_array() || input_type.get_scalar_type() != BIT { |
| 57 | return Err(runtime_error!("Clip can only be applied to bit arrays")); |
| 58 | } |
| 59 | let input_shape = input_type.get_shape(); |
| 60 | let num_bits = input_shape[input_shape.len() - 1]; |
| 61 | if self.k >= num_bits - 1 { |
| 62 | return Err(runtime_error!( |
| 63 | "Clip(k) can be applied only whenever k <= num_bits - 2" |
| 64 | )); |
| 65 | } |
| 66 | let bit_type = if input_shape.len() == 1 { |
| 67 | scalar_type(BIT) |
| 68 | } else { |
| 69 | array_type(input_shape[0..input_shape.len() - 1].to_vec(), BIT) |
| 70 | }; |
| 71 | let aux_or_graph = context.create_graph()?; |
| 72 | let state = aux_or_graph.input(bit_type.clone())?; |
| 73 | let input = aux_or_graph.input(bit_type.clone())?; |
| 74 | let output_state = |
| 75 | aux_or_graph.custom_op(CustomOperation::new(Or {}), vec![state, input])?; |
| 76 | let empty = aux_or_graph.create_tuple(vec![])?; |
| 77 | let output = aux_or_graph.create_tuple(vec![output_state, empty])?; |
| 78 | aux_or_graph.set_output_node(output)?; |
| 79 | aux_or_graph.add_annotation(GraphAnnotation::AssociativeOperation)?; |
| 80 | aux_or_graph.finalize()?; |
| 81 | let g = context.create_graph()?; |
| 82 | let input = g.input(input_type)?; |
| 83 | let input_bits = pull_out_bits(input)?; |
| 84 | let is_negative = input_bits.get(vec![num_bits - 1])?; |
| 85 | let zero_bit = g.zeros(bit_type.clone())?; |
| 86 | let one_bit = g.ones(bit_type.clone())?; |
| 87 | let top_bits = input_bits |
| 88 | .get_slice(vec![SliceElement::SubArray( |
| 89 | Some(self.k as i64), |
| 90 | None, |
| 91 | None, |
| 92 | )])? |
| 93 | .array_to_vector()?; |
| 94 | let is_large_or_negative = g |
| 95 | .iterate(aux_or_graph, zero_bit.clone(), top_bits)? |
| 96 | .tuple_get(0)?; |
| 97 | // clipped_value = if is_negative then 0, else 2^k |
| 98 | // obtained by concatenating a bunch of zeros, |
| 99 | // zero or one, then bunch of zeros again |
| 100 | let clipped_value = g |
| 101 | .create_tuple(vec![ |
| 102 | zero_bit.repeat(self.k)?, |
| 103 | g.custom_op( |
| 104 | CustomOperation::new(Mux {}), |
| 105 | vec![is_negative, zero_bit.clone(), one_bit], |
| 106 | )?, |
| 107 | zero_bit.repeat(num_bits - self.k - 1)?, |
| 108 | ])? |
nothing calls this directly
no test coverage detected