| 11 | namespace layers { |
| 12 | |
| 13 | void zero_first_timestep(StorageView& x, dim_t step) { |
| 14 | if (step == 0) { |
| 15 | x.zero(); |
| 16 | } else if (step < 0) { |
| 17 | // TODO: a more direct way to set the first timestep to 0. |
| 18 | const auto dtype = x.dtype(); |
| 19 | const auto device = x.device(); |
| 20 | StorageView first_step(dtype, device); |
| 21 | StorageView other_steps(dtype, device); |
| 22 | ops::Split(1, {1, x.dim(1) - 1})(x, first_step, other_steps); |
| 23 | first_step.zero(); |
| 24 | ops::Concat(1)({&first_step, &other_steps}, x); |
| 25 | } |
| 26 | } |
| 27 | |
| 28 | |
| 29 | Decoder::Decoder(Device device) |