tensorclouds.nn.self_interaction¶
Classes¶
Base class for all neural network modules. |
|
Base class for all neural network modules. |
|
Base class for all neural network modules. |
|
Base class for all neural network modules. |
Module Contents¶
- class tensorclouds.nn.self_interaction.FullTensorSquareSelfInteraction¶
Bases:
flax.linen.Module
Base class for all neural network modules.
Layers and models should subclass this class.
All Flax Modules are Python 3.7 dataclasses. Since dataclasses take over
__init__
, you should instead overridesetup()
, which is automatically called to initialize the module.Modules can contain submodules, and in this way can be nested in a tree structure. Submodels can be assigned as regular attributes inside the
setup()
method.You can define arbitrary “forward pass” methods on your Module subclass. While no methods are special-cased,
__call__
is a popular choice because it allows you to use module instances as if they are functions:>>> from flax import linen as nn >>> from typing import Tuple >>> class Module(nn.Module): ... features: Tuple[int, ...] = (16, 4) ... def setup(self): ... self.dense1 = nn.Dense(self.features[0]) ... self.dense2 = nn.Dense(self.features[1]) ... def __call__(self, x): ... return self.dense2(nn.relu(self.dense1(x)))
Optionally, for more concise module implementations where submodules definitions are co-located with their usage, you can use the
compact()
wrapper.- irreps: e3nn_jax.Irreps¶
- norm: bool = True¶
- class tensorclouds.nn.self_interaction.ChannelWiseTensorSquareSelfInteraction¶
Bases:
flax.linen.Module
Base class for all neural network modules.
Layers and models should subclass this class.
All Flax Modules are Python 3.7 dataclasses. Since dataclasses take over
__init__
, you should instead overridesetup()
, which is automatically called to initialize the module.Modules can contain submodules, and in this way can be nested in a tree structure. Submodels can be assigned as regular attributes inside the
setup()
method.You can define arbitrary “forward pass” methods on your Module subclass. While no methods are special-cased,
__call__
is a popular choice because it allows you to use module instances as if they are functions:>>> from flax import linen as nn >>> from typing import Tuple >>> class Module(nn.Module): ... features: Tuple[int, ...] = (16, 4) ... def setup(self): ... self.dense1 = nn.Dense(self.features[0]) ... self.dense2 = nn.Dense(self.features[1]) ... def __call__(self, x): ... return self.dense2(nn.relu(self.dense1(x)))
Optionally, for more concise module implementations where submodules definitions are co-located with their usage, you can use the
compact()
wrapper.- irreps: e3nn_jax.Irreps¶
- norm: bool = True¶
- class tensorclouds.nn.self_interaction.SegmentedTensorSquareSelfInteraction¶
Bases:
flax.linen.Module
Base class for all neural network modules.
Layers and models should subclass this class.
All Flax Modules are Python 3.7 dataclasses. Since dataclasses take over
__init__
, you should instead overridesetup()
, which is automatically called to initialize the module.Modules can contain submodules, and in this way can be nested in a tree structure. Submodels can be assigned as regular attributes inside the
setup()
method.You can define arbitrary “forward pass” methods on your Module subclass. While no methods are special-cased,
__call__
is a popular choice because it allows you to use module instances as if they are functions:>>> from flax import linen as nn >>> from typing import Tuple >>> class Module(nn.Module): ... features: Tuple[int, ...] = (16, 4) ... def setup(self): ... self.dense1 = nn.Dense(self.features[0]) ... self.dense2 = nn.Dense(self.features[1]) ... def __call__(self, x): ... return self.dense2(nn.relu(self.dense1(x)))
Optionally, for more concise module implementations where submodules definitions are co-located with their usage, you can use the
compact()
wrapper.- irreps: e3nn_jax.Irreps¶
- norm: bool = True¶
- segment_size = 2¶
- class tensorclouds.nn.self_interaction.SelfInteraction¶
Bases:
flax.linen.Module
Base class for all neural network modules.
Layers and models should subclass this class.
All Flax Modules are Python 3.7 dataclasses. Since dataclasses take over
__init__
, you should instead overridesetup()
, which is automatically called to initialize the module.Modules can contain submodules, and in this way can be nested in a tree structure. Submodels can be assigned as regular attributes inside the
setup()
method.You can define arbitrary “forward pass” methods on your Module subclass. While no methods are special-cased,
__call__
is a popular choice because it allows you to use module instances as if they are functions:>>> from flax import linen as nn >>> from typing import Tuple >>> class Module(nn.Module): ... features: Tuple[int, ...] = (16, 4) ... def setup(self): ... self.dense1 = nn.Dense(self.features[0]) ... self.dense2 = nn.Dense(self.features[1]) ... def __call__(self, x): ... return self.dense2(nn.relu(self.dense1(x)))
Optionally, for more concise module implementations where submodules definitions are co-located with their usage, you can use the
compact()
wrapper.- irreps: e3nn_jax.Irreps¶
- irreps_out: e3nn_jax.Irreps = None¶
- depth: int = 1¶
- norm_last: bool = True¶
- base: flax.linen.Module¶