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. 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
- 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]