deepxde.nn.jax

deepxde.nn.jax.fnn module

class deepxde.nn.jax.fnn.FNN(layer_sizes: Any, activation: Any, kernel_initializer: Any, params: Any = None, _input_transform: Callable = None, _output_transform: Callable = None, parent: Union[Type[flax.linen.module.Module], Type[flax.core.scope.Scope], Type[flax.linen.module._Sentinel], None] = <flax.linen.module._Sentinel object>, name: Optional[str] = None)

Bases: deepxde.nn.jax.nn.NN

Fully-connected neural network.

name = None

params = None

parent = None

scope = None

setup()

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   class MyModule(nn.Module):
     def setup(self):
       submodule = Conv(...)
   
       # Accessing `submodule` attributes does not yet work here.
   
       # The following line invokes `self.__setattr__`, which gives
       # `submodule` the name "conv1".
       self.conv1 = submodule
   
       # Accessing `submodule` attributes or methods is now safe and
       # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

deepxde.nn.jax.nn module

class deepxde.nn.jax.nn.NN(parent: Union[Type[flax.linen.module.Module], Type[flax.core.scope.Scope], Type[flax.linen.module._Sentinel], None] = <flax.linen.module._Sentinel object>, name: Optional[str] = None)

Bases: flax.linen.module.Module

Base class for all neural network modules.

apply_feature_transform(transform)

Compute the features by appling a transform to the network inputs, i.e., features = transform(inputs). Then, outputs = network(features).

apply_output_transform(transform)

Apply a transform to the network outputs, i.e., outputs = transform(inputs, outputs).

name = None

parent = None

scope = None