def METHOD_NAME(self, x):
    term1 = 0.5 * ((jnp.linalg.norm(x, axis=-1) - 2) / 0.4) ** 2
    term2 = -0.5 * ((x[..., :1] + jnp.array([-2.0, 2.0])) / 0.6) ** 2
    pe = term1 - logsumexp(term2, axis=-1)
    return -pe