def normalize_score(score: jnp.ndarray) -> jnp.ndarray:
  """Normalize score to range [0, 1] for each sample, input shape [B T H W]."""
  assert score.ndim == 4
  sample_min = score.min(axis=(1, 2, 3), keepdims=True)
  sample_max = score.max(axis=(1, 2, 3), keepdims=True)
  normalized_score = (score - sample_min) / (sample_max - sample_min + 1e-6)
  return normalized_score
Normalize score to range [0, 1] for each sample, input shape [B T H W].
videogvt/train_lib/mask_utils.py