def METHOD_NAME(scope: Scope, weights: Array):
  del scope
  norm_dims = tuple(range(weights.ndim // 2, weights.ndim))
  log_norms = jax.scipy.special.logsumexp(
      weights, axis=norm_dims, keepdims=True
  )
  return jnp.exp(weights - log_norms)