def run_models(
    video_batch1: jnp.ndarray,
    video_batch2: jnp.ndarray,
    *,
    is_tf_function: bool = True,
    metric_functions: Optional[Dict[str, MetricFunction]] = None
) -> dict[str, jnp.ndarray]:
  """Runs the jax function/models that can support jit/pmap."""
  x1, x2 = video_batch1, video_batch2
  assert x1.shape == x2.shape and x1.ndim == 5

  if metric_functions is None:
    metric_functions = dict(
        psnr=psnr,
        ssim=ssim,
    )

  flatten_t_shape = [x1.shape[0] * x1.shape[1], *x1.shape[2:]]
  flatten_x1 = jnp.reshape(x1, flatten_t_shape)
  flatten_x2 = jnp.reshape(x2, flatten_t_shape)

  if is_tf_function:
    flatten_x1 = jax.tree_map(tf.convert_to_tensor, flatten_x1)
    flatten_x2 = jax.tree_map(tf.convert_to_tensor, flatten_x2)

  ret_dict = dict()
  for (k, func) in metric_functions.items():
    flatten_metric_scores = func(flatten_x1, flatten_x2)
    if is_tf_function:
      flatten_metric_scores = flatten_metric_scores._numpy()  # pylint: disable=protected-access  # pytype: disable=attribute-error  # jax-ndarray

    ret_dict[k] = jnp.mean(
        jnp.reshape(flatten_metric_scores, x1.shape[0:2]), -1)
  return ret_dict
Runs the jax function/models that can support jit/pmap.
videogvt/train_lib/image_quality_metrics.py