format and additionnal info in jax nn

Author: bonneted

deepxde/nn/jax/nn.py

@@ -1,10 +1,12 @@
 from flax import linen as nn
+
 from .. import regularizers
 
 class NN(nn.Module):
     """Base class for all neural network modules."""
 
     # All sub-modules should have the following variables:
+    # regularization: Any = None
     # params: Any = None
     # _input_transform: Optional[Callable] = None
     # _output_transform: Optional[Callable] = None