Backend jax supports initializers, improve FNN (#547)

Author: ZongrenZou

deepxde/model.py

@@ -269,8 +269,6 @@ def closure():
 
     def _compile_jax(self, lr, loss_fn, decay, loss_weights):
         """jax"""
-        import optax
-
         # initialize network's parameters
         # TODO: Init should move to network module, because we don't know how to init here, e.g., DeepONet has two inputs.
         #       random seed should use a random number, or be specified by users
@@ -303,7 +301,7 @@ def loss_function(params):
             )  # jax.value_and_grad seems to be slightly faster than jax.grad for function approximation
             grads = grad_fn(params)
             updates, new_opt_state = self.opt.update(grads, opt_state)
-            new_params = optax.apply_updates(params, updates)
+            new_params = optimizers.apply_updates(params, updates)
             return new_params, new_opt_state
 
         def outputs(training, inputs):

deepxde/nn/initializers.py

@@ -3,7 +3,7 @@
 import math
 
 from .. import config
-from ..backend import backend_name, tf, torch
+from ..backend import backend_name, tf, torch, jax
 
 
 class VarianceScalingStacked:
@@ -139,10 +139,24 @@ def initializer_dict_torch():
     }
 
 
+def initializer_dict_jax():
+    return {
+        "Glorot normal": jax.nn.initializers.glorot_normal(),
+        "Glorot uniform": jax.nn.initializers.glorot_uniform(),
+        "He normal": jax.nn.initializers.he_normal(),
+        "He uniform": jax.nn.initializers.he_uniform(),
+        "Lecun normal": jax.nn.initializers.lecun_normal(),
+        "Lecun uniform": jax.nn.initializers.lecun_uniform(),
+        "zeros": jax.nn.initializers.zeros,
+    }
+
+
 if backend_name in ["tensorflow.compat.v1", "tensorflow"]:
     INITIALIZER_DICT = initializer_dict_tf()
 elif backend_name == "pytorch":
     INITIALIZER_DICT = initializer_dict_torch()
+elif backend_name == "jax":
+    INITIALIZER_DICT = initializer_dict_jax()
 
 
 def get(identifier):

deepxde/nn/jax/fnn.py

@@ -1,22 +1,28 @@
-from typing import Any
+from typing import Any, Callable
 
 import jax
 from flax import linen as nn
 
 from .nn import NN
+from .. import activations
+from .. import initializers
 
 
 class FNN(NN):
     """Fully-connected neural network"""
 
-    layer_sizes: Any = None
-    activation: Any = None
-    kernel_initializer: Any = None
+    layer_sizes: Any
+    activation: Any
+    kernel_initializer: Any
+    training: bool = True
+    _input_transform: Callable = None
+    _output_transform: Callable = None
+    params: Any = None
 
     def setup(self):
-        # TODO: implement get activation, get initializer
-        self._activation = jax.nn.tanh
-        kernel_initializer = jax.nn.initializers.glorot_normal()
+        # TODO: implement get regularizer
+        self._activation = activations.get(self.activation)
+        kernel_initializer = initializers.get(self.kernel_initializer)
         initializer = jax.nn.initializers.zeros
 
         self.denses = [

deepxde/nn/jax/nn.py

@@ -1,15 +1,14 @@
-from typing import Any
-
 from flax import linen as nn
 
 
 class NN(nn.Module):
     """Base class for all neural network modules."""
 
-    training: Any = True
-    params: Any = None
-    _input_transform: Any = None
-    _output_transform: Any = None
+    # all sub-modules should have the following init-only variables:
+    # training: bool = True
+    # params: Any = None
+    # _input_transform: Optional[Callable] = None
+    # _output_transform: Optional[Callable] = None
 
     def apply_feature_transform(self, transform):
         """Compute the features by appling a transform to the network inputs, i.e.,

deepxde/optimizers/jax/__init__.py

@@ -1,3 +1,3 @@
-__all__ = ["get", "is_external_optimizer"]
+__all__ = ["get", "is_external_optimizer", "apply_updates"]
 
-from .optimizers import get, is_external_optimizer
+from .optimizers import get, is_external_optimizer, apply_updates

deepxde/optimizers/jax/optimizers.py

@@ -1,9 +1,12 @@
-__all__ = ["get", "is_external_optimizer"]
+__all__ = ["get", "is_external_optimizer", "apply_updates"]
 
 import jax
 import optax
 
 
+apply_updates = optax.apply_updates
+
+
 def is_external_optimizer(optimizer):
     # TODO: add external optimizers
     return False

examples/function/func.py

@@ -1,4 +1,4 @@
-"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch"""
+"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, jax"""
 import deepxde as dde
 import numpy as np