Satisfy stricter input requirement for `jnp.linalg.norm`

[copybara-service[bot]]

Satisfy stricter input requirement for jnp.linalg.norm

The previous version of the code caused an error under the new more restrictive API of jnp.linalg.norm (see https://github.com/google/jax/pull/12670).