def illustrate_predictions(masked_inputs: jnp.ndarray, logits: jnp.ndarray,
                           targets: jnp.ndarray,
                           config: ml_collections.ConfigDict) -> jnp.ndarray:
  """illustrate the predictions as an image."""
  latent_shape = config.transformer.latent_shape
  latent_seq_len = get_latent_seq_len(config)
  masks = masked_inputs == task_registry.SpecialTokens.MASK  # TODO(Lijun-Yu): fix
  masks_3d = jnp.reshape(masks[:, -latent_seq_len:], [-1, *latent_shape])
  tile_param = [*masks_3d.shape, 1]

  light_red = jnp.tile(jnp.array([255, 229, 231]), tile_param)  # unmasked wrong
  dark_red = jnp.tile(jnp.array([139, 0, 0]), tile_param)  # masked wrong

  light_green = jnp.tile(jnp.array([229, 255, 229]),
                         tile_param)  # unmasked right
  dark_green = jnp.tile(jnp.array([21, 71, 52]), tile_param)  # masked right

  preds = jnp.argmax(logits, axis=-1)
  correct = jnp.equal(preds, targets)
  correct = jnp.reshape(correct[:, -latent_seq_len:], masks_3d.shape)

  not_correct = jnp.logical_not(correct)
  not_masks_3d = jnp.logical_not(masks_3d)

  a = jnp.logical_and(masks_3d, correct)
  b = jnp.logical_and(masks_3d, not_correct)
  c = jnp.logical_and(not_masks_3d, correct)
  # d = jnp.logical_and(not_masks_3d, not_correct)

  result_image = light_red
  result_image = jnp.where(a[..., None], dark_green, result_image)
  result_image = jnp.where(b[..., None], dark_red, result_image)
  result_image = jnp.where(c[..., None], light_green, result_image)
  return result_image / 255.
illustrate the predictions as an image.
videogvt/train_lib/maskgvt_train_lib.py