refactor decorator

deepxde/model.py

@@ -382,6 +382,7 @@ def _compile_jax(self, lr, loss_fn, decay):
         self.opt_state = self.opt.init(self.params)
 
         @jax.jit
+        @utils.list_handler
         def outputs(params, training, inputs):
             return self.net.apply(params, inputs, training=training)
 
@@ -390,9 +391,9 @@ def outputs_losses(params, training, inputs, targets, losses_fn):
 
             # TODO: Add auxiliary vars
             def outputs_fn(inputs):
-                return self.net.apply(nn_params, inputs, training=training)
+                return outputs(nn_params, training, inputs)
 
-            outputs_ = self.net.apply(nn_params, inputs, training=training)
+            outputs_ = outputs(nn_params, training, inputs)
             # Data losses
             # We use aux so that self.data.losses is a pure function.
             aux = [outputs_fn, ext_params] if ext_params else [outputs_fn]

deepxde/nn/jax/snn.py

@@ -9,7 +9,6 @@
 from .. import initializers
 from ...utils import list_handler
 
-@list_handler
 class SPINN(NN):
     layer_sizes: Any
     activation: Any

deepxde/utils/internal.py

@@ -235,10 +235,10 @@ def mpi_scatter_from_rank0(array, drop_last=True):
 
 def list_handler(func):
     @wraps(func)
-    def wrapper(*args):
+    def wrapper(*args, **kwargs):
         inputs = args[-1]
         if isinstance(inputs, (list, tuple)):
-            results = [func(*args[:-1], input_item) for input_item in inputs]
+            results = [func(*args[:-1], input_item, **kwargs) for input_item in inputs]
             return bkd.concat(results, axis=0)
         return func(*args)