def _average_pooling_3d_in_place(
    value: jnp.ndarray,
    poolable: jnp.ndarray,
    window_size: tuple[int, int, int],
):
  """Average pooling 3D on poolable locations while keeping original shape."""
  bs, t, h, w = value.shape
  t_b, h_b, w_b = window_size
  grouped_shape = (bs, t // t_b, t_b, h // h_b, h_b, w // w_b, w_b)
  value = jnp.reshape(value, grouped_shape)
  poolable = jnp.reshape(poolable, grouped_shape)

  value_avg = jnp.mean(value, axis=(2, 4, 6), keepdims=True, where=poolable)
  tiles = (1, 1, t_b, 1, h_b, 1, w_b)
  value_avg = jnp.tile(value_avg, tiles)
  value = jnp.where(poolable, value_avg, value)

  value = jnp.reshape(value, (bs, t, h, w))
  return value
Average pooling 3D on poolable locations while keeping original shape.
videogvt/train_lib/mask_utils.py