def lm_decode(seq: jnp.ndarray,
              prefix_lengths: jnp.ndarray,
              cache: Dict[str, jnp.ndarray],
              rng: jnp.ndarray,
              variables,
              flax_model: nn.Module,
              config: ml_collections.ConfigDict) -> jnp.ndarray:
  """Fast decoding method.

  Args:
    seq: [batch, cond_len] input tokens.
    prefix_lengths: [batch,] prefix lengths.
    cache: cached variables for self-attention layers.
    rng: PRNGKey for decoding.
    variables: transformer model weights.
    flax_model: transformer model object.
    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(inputs, cache, enable_vq_codebook=jnp.array(False, bool)):
    new_variables = dict(list(variables.items()) + [('cache', cache)])
    outputs = flax_model.apply(
        new_variables,
        inputs,
        mutable=['cache'],
        deterministic=True)
    # [bs, 1, vocab_size]
    logits, new_cache = outputs
    all_valid = jnp.ones((logits.shape[-1],), logits.dtype)
    vq_codebook_valid = jnp.concatenate([
        jnp.ones((vq_codebook_size,), logits.dtype),
        jnp.zeros((logits.shape[-1] - vq_codebook_size,), logits.dtype)
    ],
                                        axis=0)
    logit_valid = jnp.where(enable_vq_codebook, vq_codebook_valid, all_valid)
    logits = logits * jnp.broadcast_to(logit_valid, logits.shape) - 1000 * (
        1 - jnp.broadcast_to(logit_valid, logits.shape))
    return logits, new_cache['cache']
  output_indices = lm_sample_decode.decode(
      seq,
      prefix_lengths,
      cache,
      decode_fn,
      rng,
      sampling_topk=config.sampling.sampling_topk,
      sampling_topp=config.sampling.sampling_topp,
      sampling_temperature=config.sampling.sampling_temperature)
  output_indices = output_indices[:, -latent_seq_len:]
  return output_indices
Fast decoding method.

  Args:
    seq: [batch, cond_len] input tokens.
    prefix_lengths: [batch,] prefix lengths.
    cache: cached variables for self-attention layers.
    rng: PRNGKey for decoding.
    variables: transformer model weights.
    flax_model: transformer model object.
    config: model config.

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