Backend JAX: Add norm (#1950)

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, axis=axis, keepdims=keepdims)
+
+
 def zeros(shape, dtype):
     return jnp.zeros(shape, dtype=dtype)