neuralop.layers.embeddings.GridEmbeddingND
- class neuralop.layers.embeddings.GridEmbeddingND(in_channels: int, dim: int = 2, grid_boundaries=[[0, 1], [0, 1]])[source]
GridEmbeddingND applies a simple positional embedding as a regular ND grid.
It expects inputs of shape (batch, channels, d_1, …, d_n)
- Parameters:
- in_channelsint
number of channels in input
- dimint
dimensions of positional encoding to apply
- grid_boundarieslist, optional
coordinate boundaries of input grid along each dim, by default [[0, 1], [0, 1]]
- 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.
- devicestr
Device to load data (e.g. ‘cpu’ or ‘cuda:*’).
- dtypestr
dtype to encode data.
- Returns:
- torch.tensor
output grids to concatenate
- forward(data, batched=True)[source]
- Parameters:
- data: torch.Tensor
assumes shape (batch (optional), channels, x_1, x_2, …x_n)
- batched: bool
whether data has a batch dim