emmi_inference.models.modules.continuous_sincos_embed

Classes

ContinuousSincosEmbed

Embedding layer for continuous coordinates using sine and cosine functions as used in transformers.

Module Contents

class emmi_inference.models.modules.continuous_sincos_embed.ContinuousSincosEmbed(dim, ndim, max_wavelength=10000, assert_positive=True)

Bases: torch.nn.Module

Embedding layer for continuous coordinates using sine and cosine functions as used in transformers. This implementation is able to deal with arbitrary coordinate dimensions (e.g., 2D and 3D coordinate systems).

Parameters:

  • dim (int) – Dimensionality of the embedded input coordinates.

  • ndim (int) – Number of dimensions of the input domain.

  • max_wavelength (int) – Max length. Defaults to 10000.

  • assert_positive (bool) – If true, assert if all input coordiantes are positive. Defaults to True.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

dim
ndim
ndim_padding
sincos_padding
max_wavelength = 10000
padding
assert_positive = True
forward(coords)

Forward method of the ContinuousSincosEmbed layer.

Parameters:

coords (torch.Tensor) – Tensor of coordinates. The shape of the tensor should be (batch size, number of points, coordinate dimension) or (number of points, coordinate dimension).

Returns:

Tensor with embedded coordinates.

Return type:

torch.Tensor