def METHOD_NAME(batch_shape):
    """Get parameterized Normal with given batch shape."""
    loc = jnp.zeros(batch_shape)
    scale = jnp.ones(batch_shape)
    normal = dist.Normal(loc=loc, scale=scale)
    return normal