refactor 3D Exception to gradients_reverse.py

Author: Jerry-Jzy

deepxde/gradients/gradients.py

@@ -4,7 +4,6 @@
 
 from . import gradients_forward
 from . import gradients_reverse
-from .. import backend as bkd
 from .. import config
 
 
@@ -37,10 +36,6 @@ def jacobian(ys, xs, i=None, j=None):
         (batch_size_out, batch_size, 1).
     """
     if config.autodiff == "reverse":
-        if bkd.ndim(ys) == 3:
-            raise NotImplementedError(
-                "Reverse-mode autodiff doesn't support 3D output"
-            )
         return gradients_reverse.jacobian(ys, xs, i=i, j=j)
     if config.autodiff == "forward":
         return gradients_forward.jacobian(ys, xs, i=i, j=j)
@@ -72,10 +67,6 @@ def hessian(ys, xs, component=0, i=0, j=0):
         the output shape is (batch_size_out, batch_size, 1).
     """
     if config.autodiff == "reverse":
-        if bkd.ndim(ys) == 3:
-            raise NotImplementedError(
-                "Reverse-mode autodiff doesn't support 3D output"
-            )
         return gradients_reverse.hessian(ys, xs, component=component, i=i, j=j)
     if config.autodiff == "forward":
         return gradients_forward.hessian(ys, xs, component=component, i=i, j=j)

deepxde/gradients/gradients_reverse.py

@@ -3,7 +3,7 @@
 __all__ = ["hessian", "jacobian"]
 
 from .jacobian import Jacobian, Jacobians
-from ..backend import backend_name, tf, torch, jax, paddle
+from ..backend import backend_name, bkd, tf, torch, jax, paddle
 
 
 class JacobianReverse(Jacobian):
@@ -21,11 +21,11 @@ def __call__(self, i=None, j=None):
         # Compute J[i, :]
         if i not in self.J:
             if backend_name in ["tensorflow.compat.v1", "tensorflow"]:
-                y = self.ys[..., i : i + 1] if self.dim_y > 1 else self.ys
+                y = self.ys[:, i : i + 1] if self.dim_y > 1 else self.ys
                 self.J[i] = tf.gradients(y, self.xs)[0]
             elif backend_name == "pytorch":
                 # TODO: retain_graph=True has memory leak?
-                y = self.ys[..., i : i + 1] if self.dim_y > 1 else self.ys
+                y = self.ys[:, i : i + 1] if self.dim_y > 1 else self.ys
                 self.J[i] = torch.autograd.grad(
                     y, self.xs, grad_outputs=torch.ones_like(y), create_graph=True
                 )[0]
@@ -43,7 +43,7 @@ def __call__(self, i=None, j=None):
                 grad_fn = jax.grad(lambda x: self.ys[1](x)[i])
                 self.J[i] = (jax.vmap(grad_fn)(self.xs), grad_fn)
             elif backend_name == "paddle":
-                y = self.ys[..., i : i + 1] if self.dim_y > 1 else self.ys
+                y = self.ys[:, i : i + 1] if self.dim_y > 1 else self.ys
                 self.J[i] = paddle.grad(y, self.xs, create_graph=True)[0]
 
         if j is None or self.dim_x == 1:
@@ -57,19 +57,23 @@ def __call__(self, i=None, j=None):
                 "pytorch",
                 "paddle",
             ]:
-                self.J[i, j] = self.J[i][..., j : j + 1]
+                self.J[i, j] = self.J[i][:, j : j + 1]
             elif backend_name == "jax":
                 # In backend jax, a tuple of a jax array and a callable is returned, so
                 # that it is consistent with the argument, which is also a tuple. This
                 # is useful for further computation, e.g., Hessian.
                 self.J[i, j] = (
-                    self.J[i][0][..., j : j + 1],
+                    self.J[i][0][:, j : j + 1],
                     lambda x: self.J[i][1](x)[j : j + 1],
                 )
         return self.J[i, j]
 
 
 def jacobian(ys, xs, i=None, j=None):
+    if bkd.ndim(ys) == 3:
+        raise NotImplementedError(
+            "Reverse-mode autodiff doesn't support 3D output"
+        )
     return jacobian._Jacobians(ys, xs, i=i, j=j)
 
 
@@ -137,6 +141,10 @@ def clear(self):
 
 
 def hessian(ys, xs, component=0, i=0, j=0):
+    if bkd.ndim(ys) == 3:
+        raise NotImplementedError(
+            "Reverse-mode autodiff doesn't support 3D output"
+        )
     return hessian._Hessians(ys, xs, component=component, i=i, j=j)