def METHOD_NAME(i, x):
    mean, log_scale = self.arn(x)
    inverse_scale = jnp.exp(
        -_clamp_preserve_gradients(
            log_scale, min=self.log_scale_min_clip, max=self.log_scale_max_clip
        )
    )
    x = (y - mean) * inverse_scale
    return x