p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] sample.py (python)
def _sampling_wrapper(module, batch):
    return module.sample(batch, args.seq_len, args.maskgit_steps, args.temperature, args.sample_argmax)

# --- Define autoregressive sampling loop ---
def _autoreg_sample(rng, video_batch, action_batch):
    vid = video_batch[:, : args.start_frame + 1]
    sampling_fn = jax.jit(nn.apply(_sampling_wrapper, genie)) 
    rng, _rng = jax.random.split(rng)
    batch = dict(videos=vid, latent_actions=action_batch, rng=_rng)
    generated_vid = sampling_fn(
        params,
        batch
    )
    return generated_vid