add norm for jax backend

Author: bonneted

deepxde/backend/jax/tensor.py

@@ -165,6 +165,10 @@ def reduce_max(input_tensor):
     return jnp.max(input_tensor)
 
 
+def norm(tensor, ord=None, axis=None, keepdims=False):
+    return jnp.linalg.norm(tensor, ord=ord, dim=axis, keepdim=keepdims)
+
+
 def zeros(shape, dtype):
     return jnp.zeros(shape, dtype=dtype)