fix param name

Author: bonneted

deepxde/backend/jax/tensor.py

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