def resize_bilinear(img, new_size):
  """Bilinear resizing matching the behavior of TF 1.x align_corners=False.

  See https://github.com/google/jax/issues/862#issuecomment-567674762

  Args:
    img: image
    new_size: tuple of integers representing the new height and width

  Returns:
    resized image
  """
  in_rows, in_cols = img.shape[-3:-1]
  new_rows, new_cols = new_size
  rows = jnp.arange(new_rows, dtype=jnp.float32) * (in_rows / new_rows)
  cols = jnp.arange(new_cols, dtype=jnp.float32) * (in_cols / new_cols)
  rows2, cols2 = jnp.meshgrid(rows, cols, indexing='ij')
  img_resize_vec = _interpolate_bilinear(img, rows2.flatten(), cols2.flatten())
  img_resize = img_resize_vec.reshape(
      img.shape[:-3] + (len(rows), len(cols)) + img.shape[-1:])
  return img_resize
Bilinear resizing matching the behavior of TF 1.x align_corners=False.

  See https://github.com/google/jax/issues/862#issuecomment-567674762

  Args:
    img: image
    new_size: tuple of integers representing the new height and width

  Returns:
    resized image
  
videogvt/train_lib/metrics_lib.py