(self, input_data: ArgData, fixed_params, rng)
| 111 | self.unfolded_input = None |
| 112 | |
| 113 | def setup(self, input_data: ArgData, fixed_params, rng): |
| 114 | self.input_data = input_data |
| 115 | self.fixed_params = fixed_params |
| 116 | self.rng = rng |
| 117 | |
| 118 | def unfold_output(self, batches): |
| 119 | num_expand = len(self.input_data.desc.expandable_prefix) |