def go():
    # Note that value for learning_rate is arbitrary, but we pass it here to
    # warm up the jit cache (since we are passing a learning rate at training
    # time).
    res = parallel_train_step(
        replicated_optimizer,
        batched_example,
        static_batch_metadata,
        loss_fn,
        max_global_norm=max_global_norm,
        learning_rate=0.0)
    jax.tree_map(lambda x: x.block_until_ready(), res)