deepxde.nn.jax
deepxde.nn.jax.fnn module
- class deepxde.nn.jax.fnn.FNN(layer_sizes: ~typing.Any, activation: ~typing.Any, kernel_initializer: ~typing.Any, params: ~typing.Any = None, _input_transform: ~typing.Callable = None, _output_transform: ~typing.Callable = None, parent: ~flax.linen.module.Module | ~flax.core.scope.Scope | ~flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)[source]
Bases:
NN
Fully-connected neural network.
- activation: Any
- kernel_initializer: Any
- layer_sizes: Any
- name: str | None = None
- params: Any = None
- parent: Module | Scope | _Sentinel | None = None
- scope: Scope | None = None
- setup()[source]
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 asetup
-defined attribute onself
is accessed.This can happen in three cases:
Immediately when invoking
apply()
,init()
orinit_and_output()
.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 = nn.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.
Once a module is constructed inside a method wrapped with
compact()
, immediately before another method is called orsetup
defined attribute is accessed.
- class deepxde.nn.jax.fnn.PFNN(layer_sizes: ~typing.Any, activation: ~typing.Any, kernel_initializer: ~typing.Any, params: ~typing.Any = None, _input_transform: ~typing.Callable = None, _output_transform: ~typing.Callable = None, parent: ~flax.linen.module.Module | ~flax.core.scope.Scope | ~flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)[source]
Bases:
NN
Parallel fully-connected network that uses independent sub-networks for each network output.
- Parameters:
layer_sizes – A nested list that defines the architecture of the neural network (how the layers are connected). If layer_sizes[i] is an int, it represents one layer shared by all the outputs; if layer_sizes[i] is a list, it represents len(layer_sizes[i]) sub-layers, each of which is exclusively used by one output. Every layer_sizes[i] list must have the same length (= number of subnetworks). If the last element of layer_sizes is an int preceded by a list, it must be equal to the number of subnetworks: all subnetworks have an output size of 1 and are then concatenated. If the last element is a list, it specifies the output size for each subnetwork before concatenation.
- activation: Any
- kernel_initializer: Any
- layer_sizes: Any
- name: str | None = None
- params: Any = None
- parent: Module | Scope | _Sentinel | None = None
- scope: Scope | None = None
- setup()[source]
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 asetup
-defined attribute onself
is accessed.This can happen in three cases:
Immediately when invoking
apply()
,init()
orinit_and_output()
.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 = nn.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.
Once a module is constructed inside a method wrapped with
compact()
, immediately before another method is called orsetup
defined attribute is accessed.
deepxde.nn.jax.nn module
- class deepxde.nn.jax.nn.NN(parent: ~flax.linen.module.Module | ~flax.core.scope.Scope | ~flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)[source]
Bases:
Module
Base class for all neural network modules.
- apply_feature_transform(transform)[source]
Compute the features by appling a transform to the network inputs, i.e., features = transform(inputs). Then, outputs = network(features).
- apply_output_transform(transform)[source]
Apply a transform to the network outputs, i.e., outputs = transform(inputs, outputs).
- name: str | None = None
- parent: Module | Scope | _Sentinel | None = None
- scope: Scope | None = None