neuralop.layers.embeddings.GridEmbeddingND

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

A positional embedding as a regular ND grid

Attributes:

out_channels

Methods

forward(data[, batched])
grid(spatial_dims, device, dtype)	grid generates ND grid needed for pos encoding and caches the grid associated with MRU resolution

grid(spatial_dims: Size, device: str, dtype: dtype)[source]

grid generates ND 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]