Tensorflow 2.x backend: multiple outputs extension of DeepONet

deepxde/data/triple.py

@@ -29,6 +29,11 @@ def __init__(self, X_train, y_train, X_test, y_test):
         self.train_sampler = BatchSampler(len(self.train_y), shuffle=True)
 
     def losses(self, targets, outputs, loss_fn, inputs, model, aux=None):
+        if isinstance(loss_fn, list):
+            losses = []
+            for fn in loss_fn:
+                losses.append(fn(targets, outputs))
+            return losses
         return loss_fn(targets, outputs)
 
     def train_next_batch(self, batch_size=None):
@@ -74,14 +79,21 @@ def __init__(self, X_train, y_train, X_test, y_test):
         self.trunk_sampler = BatchSampler(len(X_train[1]), shuffle=True)
 
     def losses(self, targets, outputs, loss_fn, inputs, model, aux=None):
+        if isinstance(loss_fn, list):
+            losses = []
+            for fn in loss_fn:
+                losses.append(fn(targets, outputs))
+            return losses
         return loss_fn(targets, outputs)
 
     def train_next_batch(self, batch_size=None):
         if batch_size is None:
             return self.train_x, self.train_y
         if not isinstance(batch_size, (tuple, list)):
             indices = self.branch_sampler.get_next(batch_size)
-            return (self.train_x[0][indices], self.train_x[1]), self.train_y[indices]
+            return (self.train_x[0][indices], self.train_x[1]), self.train_y[
+                indices
+            ]
         indices_branch = self.branch_sampler.get_next(batch_size[0])
         indices_trunk = self.trunk_sampler.get_next(batch_size[1])
         return (

deepxde/nn/tensorflow/__init__.py

@@ -1,6 +1,13 @@
 """Package for tensorflow NN modules."""
 
-__all__ = ["DeepONet", "DeepONetCartesianProd", "FNN", "NN", "PFNN", "PODDeepONet"]
+__all__ = [
+    "DeepONet",
+    "DeepONetCartesianProd",
+    "FNN",
+    "NN",
+    "PFNN",
+    "PODDeepONet",
+]
 
 from .deeponet import DeepONet, DeepONetCartesianProd, PODDeepONet
 from .fnn import FNN, PFNN