neuralop.layers.embeddings.GridEmbedding2D

class neuralop.layers.embeddings.GridEmbedding2D(in_channels: int, grid_boundaries=[[0, 1], [0, 1]])[source]

A simple positional embedding as a regular 2D grid

Attributes:
out_channels

Methods

forward(data[, batched])

Define the computation performed at every call.

grid(spatial_dims, device, dtype)

grid generates 2D grid needed for pos encoding and caches the grid associated with MRU resolution

grid(spatial_dims, device, dtype)[source]

grid generates 2D grid needed for pos encoding and caches the grid associated with MRU resolution

Parameters:
spatial_dimstorch.size

sizes of spatial resolution

deviceliteral ‘cpu’ or ‘cuda:*

where to load data

dtypestr

dtype to encode data

Returns:
torch.tensor

output grids to concatenate

forward(data, batched=True)[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.