def fast_decode(batch: Batch, prefix_token_dict: jnp.ndarray,
                mask_fn: mask_utils.MaskFn, rng: jnp.ndarray, variables,
                flax_model: nn.Module,
                inference_method: Optional[Callable[..., Any]],
                config: ml_collections.ConfigDict) -> jnp.ndarray:
  """Fast decoding method.

  Args:
    batch: Dictionary where batch['inputs'].shape=[batch, T, H, W] is the full
      tokens, and optionally batch['cond_inputs'] and batch['cond_mask'] in a
      broadcastable shape to batch['inputs'] are the condition tokens and masks.
    prefix_token_dict: Dictionary where prefix_token_dict['masked_inputs'].shape
      =[batch, cond_len] is the input prefix tokens and
      prefix_token_dict['segment_ids'] is segment ids in the same shape.
    mask_fn: mask function.
    rng: PRNGKey for decoding.
    variables: transformer model weights.
    flax_model: transformer model object.
    inference_method: optional method to inference transformer.
    config: model config.

  Returns:
    [batch, seq_len] decoded token sequence.
  """
  vq_codebook_size = get_vq_codebook_size(config)
  latent_seq_len = get_latent_seq_len(config)

  def decode_fn(cur_inputs: jnp.ndarray,
                segment_ids: jnp.ndarray,
                cur_uncond_inputs: jnp.ndarray | None = None) -> jnp.ndarray:
    del cur_uncond_inputs
    masked_inputs = mask_utils.concat_tokens(prefix_token_dict['masked_inputs'],
                                             cur_inputs)
    segment_ids = mask_utils.concat_tokens(prefix_token_dict['segment_ids'],
                                           segment_ids)
    model_inputs = [masked_inputs, segment_ids]
    if 'lowres_tokens' in batch:
      model_inputs.append(batch['lowres_tokens'])

    logits = flax_model.apply(
        variables,
        *model_inputs,
        method=inference_method,
        deterministic=True)
    logits = logits[:, -latent_seq_len:, :vq_codebook_size]
    logits = logits.reshape(*cur_inputs.shape, vq_codebook_size)
    return logits

  def decode_fn_with_guidance(cur_inputs: jnp.ndarray, segment_ids: jnp.ndarray,
                              cur_uncond_inputs: jnp.ndarray,
                              guidance_scale: float) -> jnp.ndarray:
    masked_inputs = mask_utils.concat_tokens(prefix_token_dict['masked_inputs'],
                                             cur_inputs)
    segment_ids = mask_utils.concat_tokens(prefix_token_dict['segment_ids'],
                                           segment_ids)
    model_inputs = [masked_inputs, segment_ids]
    if 'lowres_tokens' in batch:
      model_inputs.append(batch['lowres_tokens'])
    logits = flax_model.apply(
        variables,
        *model_inputs,
        method=inference_method,
        deterministic=True)
    logits = logits[:, -latent_seq_len:, :vq_codebook_size]
    uncond_inputs = mask_utils.concat_tokens(
        prefix_token_dict['uncond_masked_inputs'], cur_uncond_inputs
    )
    uncond_logits = flax_model.apply(
        variables,
        uncond_inputs,
        segment_ids,
        method=inference_method,
        deterministic=True)
    uncond_logits = uncond_logits[:, -latent_seq_len:, :vq_codebook_size]
    logits = logits + guidance_scale * (logits - uncond_logits)
    logits = logits.reshape(*cur_inputs.shape, vq_codebook_size)
    return logits

  def mask_fn_with_cond(cur_inputs: jnp.ndarray, *args,
                        **kwargs) -> mask_utils.TokenDict:
    cur_batch = {'inputs': cur_inputs}
    if 'cond_inputs' in batch:
      cur_batch['cond_inputs'] = batch['cond_inputs']
    if 'cond_mask' in batch:
      cur_batch['cond_mask'] = batch['cond_mask']
    return mask_fn(cur_batch, *args, **kwargs)

  guidance_scale = config.sampling.get('guidance_scale', 0.0)
  output_indices = parallel_decode.decode(
      batch['inputs'],
      decode_fn if guidance_scale == 0.0 else functools.partial(
          decode_fn_with_guidance, guidance_scale=guidance_scale),
      mask_fn_with_cond,
      rng,
      num_iter=config.sampling.mask_bins,
      sampling_topk=config.sampling.get('sampling_topk', 0),
      sampling_topp=config.sampling.get('sampling_topp', 0.),
      mask_temperature=config.sampling.choice_temperature,
      mask_scheduling_method=config.sampling.mask_scheduling_method)
  output_indices = jnp.array(output_indices[:, -1], dtype=jnp.int32)
  return output_indices
Fast decoding method.

  Args:
    batch: Dictionary where batch['inputs'].shape=[batch, T, H, W] is the full
      tokens, and optionally batch['cond_inputs'] and batch['cond_mask'] in a
      broadcastable shape to batch['inputs'] are the condition tokens and masks.
    prefix_token_dict: Dictionary where prefix_token_dict['masked_inputs'].shape
      =[batch, cond_len] is the input prefix tokens and
      prefix_token_dict['segment_ids'] is segment ids in the same shape.
    mask_fn: mask function.
    rng: PRNGKey for decoding.
    variables: transformer model weights.
    flax_model: transformer model object.
    inference_method: optional method to inference transformer.
    config: model config.

  Returns:
    [batch, seq_len] decoded token sequence.
  
videogvt/train_lib/maskgvt_train_lib.py