move regularizer to base nn class

Author: bonneted

deepxde/nn/jax/fnn.py

@@ -7,8 +7,6 @@
 from .nn import NN
 from .. import activations
 from .. import initializers
-from .. import regularizers
-
 
 class FNN(NN):
     """Fully-connected neural network."""
@@ -22,11 +20,6 @@ class FNN(NN):
     _input_transform: Callable = None
     _output_transform: Callable = None
 
-    @property
-    def regularizer(self):
-        """Dynamically compute and return the regularizer function based on regularization."""
-        return regularizers.get(self.regularization)
-    
     def setup(self):
         # TODO: implement get regularizer
         if isinstance(self.activation, list):
@@ -91,11 +84,6 @@ class PFNN(NN):
     _input_transform: Callable = None
     _output_transform: Callable = None
 
-    @property
-    def regularizer(self):
-        """Dynamically compute and return the regularizer function based on regularization."""
-        return regularizers.get(self.regularization)
-
     def setup(self):
         if len(self.layer_sizes) <= 1:
             raise ValueError("must specify input and output sizes")

deepxde/nn/jax/nn.py

@@ -1,5 +1,5 @@
 from flax import linen as nn
-
+from .. import regularizers
 
 class NN(nn.Module):
     """Base class for all neural network modules."""
@@ -8,7 +8,12 @@ class NN(nn.Module):
     # params: Any = None
     # _input_transform: Optional[Callable] = None
     # _output_transform: Optional[Callable] = None
-
+    
+    @property
+    def regularizer(self):
+        """Dynamically compute and return the regularizer function based on regularization."""
+        return regularizers.get(self.regularization)
+    
     def apply_feature_transform(self, transform):
         """Compute the features by appling a transform to the network inputs, i.e.,
         features = transform(inputs). Then, outputs = network(features).