Expands the `arg_data` to match the sequence shape of input_data.
(input_data: ArgData, arg_data: ArgData)
| 252 | |
| 253 | |
| 254 | def expand_arg_input(input_data: ArgData, arg_data: ArgData): |
| 255 | """ |
| 256 | Expands the `arg_data` to match the sequence shape of input_data. |
| 257 | """ |
| 258 | assert arg_data.desc.expandable_prefix in ["F", ""] |
| 259 | assert len(input_data.data) == len(arg_data.data) |
| 260 | arg_has_frames = arg_data.desc.expandable_prefix == "F" |
| 261 | return [ |
| 262 | expand_arg(input_data.desc.expandable_prefix, arg_has_frames, input_batch, arg_batch) |
| 263 | for input_batch, arg_batch in zip(input_data.data, arg_data.data) |
| 264 | ] |
| 265 | |
| 266 | |
| 267 | def _test_seq_input(num_iters, operator_fn, fixed_params, input_params, input_data: ArgData, rng): |
no test coverage detected