Backend JAX supports model.predict (#1594)

Author: ZongrenZou

deepxde/model.py

@@ -935,6 +935,22 @@ def op(inputs):
             # Clear cached Jacobians and Hessians.
             grad.clear()
             y = utils.to_numpy(y)
+        elif backend_name == "jax":
+            if utils.get_num_args(operator) == 2:
+
+                @jax.jit
+                def op(inputs):
+                    y_fn = lambda _x: self.net.apply(self.net.params, _x)
+                    return operator(inputs, (y_fn(inputs), y_fn))
+
+            elif utils.get_num_args(operator) == 3:
+                # TODO: JAX backend Implementation of Auxiliary variables.
+                raise NotImplementedError(
+                    "Model.predict() with auxiliary variable hasn't been implemented "
+                    "for backend jax."
+                )
+            y = op(x)
+            y = utils.to_numpy(y)
         elif backend_name == "paddle":
             self.net.eval()
             inputs = paddle.to_tensor(x, stop_gradient=False)