Compute Euclidean norm using custom jvp implementation.
Here we use a custom jvp implementation for the norm that does not yield
NaN gradients when differentiating the norm of (x-x), but defaults
instead to zero, using a custom_jvp rule.
- Parameters
-
- Return type
float
