(scan_state: Any)
| 155 | return i < num_steps |
| 156 | |
| 157 | def loop_fn(scan_state: Any) -> Tuple[Any, Any, Any, Any]: |
| 158 | (dstate, input_token, i, _) = scan_state |
| 159 | |
| 160 | (logits, dstate, _) = self.decoder( |
| 161 | input_tokens=input_token, |
| 162 | target_tokens=None, |
| 163 | start_of_sequence=no_start_of_seq, |
| 164 | decoder_state=dstate, |
| 165 | ) |
| 166 | |
| 167 | logits = logits / temperature |
| 168 | output_token = jax.lax.dynamic_slice_in_dim(target_tokens, i, 1, axis=1) |
| 169 | |
| 170 | return (dstate, output_token, i + 1, logits) |
| 171 | |
| 172 | # Scan over the sequence length. |
| 173 | dummy_logits = jnp.zeros((batch_size, 1, 1024)) |
nothing calls this directly
no outgoing calls
no test coverage detected