Activations#

class physicsnemo.nn.module.activations.CappedGELU(cap_value=1.0, **kwargs)[source]#

Bases: Module

Implements a GELU with capped maximum value.

Example

>>> capped_gelu_func = physicsnemo.nn.CappedGELU()
>>> input = torch.Tensor([[-2,-1],[0,1],[2,3]])
>>> capped_gelu_func(input)
tensor([[-0.0455, -0.1587],
        [ 0.0000,  0.8413],
        [ 1.0000,  1.0000]])

forward(inputs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.nn.module.activations.CappedLeakyReLU(cap_value=1.0, **kwargs)[source]#

Bases: Module

Implements a ReLU with capped maximum value.

Example

>>> capped_leakyReLU_func = physicsnemo.nn.CappedLeakyReLU()
>>> input = torch.Tensor([[-2,-1],[0,1],[2,3]])
>>> capped_leakyReLU_func(input)
tensor([[-0.0200, -0.0100],
        [ 0.0000,  1.0000],
        [ 1.0000,  1.0000]])

forward(inputs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.nn.module.activations.Identity(*args: Any, **kwargs: Any)[source]#

Bases: Module

Identity activation function

Dummy function for removing activations from a model

Example

>>> idnt_func = physicsnemo.nn.Identity()
>>> input = torch.randn(2, 2)
>>> output = idnt_func(input)
>>> torch.allclose(input, output)
True

forward(x: Tensor) → Tensor[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.nn.module.activations.SquarePlus[source]#

Bases: Module

Squareplus activation

Note

Reference: arXiv preprint arXiv:2112.11687

Example

>>> sqr_func = physicsnemo.nn.SquarePlus()
>>> input = torch.Tensor([[1,2],[3,4]])
>>> sqr_func(input)
tensor([[1.6180, 2.4142],
        [3.3028, 4.2361]])

forward(x: Tensor) → Tensor[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.nn.module.activations.Stan(out_features: int = 1)[source]#

Bases: Module

Self-scalable Tanh (Stan) for 1D Tensors

Parameters:: out_features (int, optional) – Number of features, by default 1

Note

References: Gnanasambandam, Raghav and Shen, Bo and Chung, Jihoon and Yue, Xubo and others. Self-scalable Tanh (Stan): Faster Convergence and Better Generalization in Physics-informed Neural Networks. arXiv preprint arXiv:2204.12589, 2022.

Example

>>> stan_func = physicsnemo.nn.Stan(out_features=1)
>>> input = torch.Tensor([[0],[1],[2]])
>>> stan_func(input)
tensor([[0.0000],
        [1.5232],
        [2.8921]], grad_fn=<MulBackward0>)

forward(x: Tensor) → Tensor[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

physicsnemo.nn.module.activations.get_activation(activation: str) → Module[source]#

Returns an activation function given a string

Parameters:: activation (str) – String identifier for the desired activation function
Return type:: Activation function
Raises:: KeyError – If the specified activation function is not found in the dictionary

class physicsnemo.nn.module.fused_silu.FusedSiLU[source]#

Bases: Function

Placeholder for when nvfuser is not available.

class physicsnemo.nn.module.gumbel_softmax.GumbelSoftmax(tau: float = 1.0, learnable: bool = False)[source]#

Bases: Module

Gumbel-Softmax module for differentiable categorical sampling.

This module wraps the gumbel_softmax() function as an nn.Module, allowing it to be used as a layer in neural network architectures.

The Gumbel-Softmax trick provides a differentiable approximation to sampling from a categorical distribution, enabling end-to-end training of models with discrete latent variables.

Parameters:

tau (float, optional, default=1.0) – Initial temperature parameter. Lower values make the distribution more concentrated (closer to one-hot). Can be modified after initialization.
learnable (bool, optional, default=False) – If True, the temperature parameter is registered as a learnable nn.Parameter. If False, it is a fixed buffer.

Examples

>>> import torch
>>> gs = GumbelSoftmax(tau=0.5)
>>> logits = torch.randn(2, 10)  # batch_size=2, num_categories=10
>>> probs = gs(logits)
>>> probs.shape
torch.Size([2, 10])
>>> torch.allclose(probs.sum(dim=-1), torch.ones(2))  # Each row sums to 1
True

>>> # With learnable temperature
>>> gs_learnable = GumbelSoftmax(tau=1.0, learnable=True)
>>> gs_learnable.tau.requires_grad
True