add norm for jax backend#1950

lululxvi merged 3 commits intolululxvi:masterlululxvi/deepxde:masterfrom

bonneted:masterbonneted/deepxde:masterCopy head branch name to clipboard