update

MaxText/attentions_test.py

@@ -130,12 +130,12 @@ def main(argv: Sequence[str]) -> None:
   # empty variable states in pure flash attention
   vars = {}
   jitted = jax.jit(nn.apply(test_tpu_flash_attention, attention_op))
+  query, key, value, decoder_segment_ids = next(train_iter)
   for step in np.arange(start_step, config.steps):
     if step == first_profiling_step:
       prof.activate()
     with jax.profiler.StepTraceAnnotation("train", step_num=step):
       with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
-        query, key, value, decoder_segment_ids = next(train_iter)
         out = jitted(vars, query, key, value, decoder_segment_ids)
     # jax.block_until_ready(out)
     # pdb.set_trace()