Skip to content

inro.models

The architectures of the model zoo, all re-exported at the top level (inro.SIREN, inro.InstantNGP, ...). FUTON has its own page.

MLP

mlp

Multi-layer perceptrons with standard or custom hidden layers.

ReLULayer

ReLULayer(in_features: int, out_features: int)

Linear projection followed by ReLU.

Parameters:

  • in_features (int) –

    Input width.

  • out_features (int) –

    Output width.

Shape
  • Input: (..., in_features).
  • Output: (..., out_features).
Source code in inro/models/mlp.py
def __init__(self, in_features: int, out_features: int) -> None:
    super().__init__()
    self.in_features = in_features
    self.linear = nn.Linear(in_features, out_features)

MLP

MLP(in_features: int, out_features: int, hidden_features: int | None = None, hidden_layers: int = 2, activation: Callable[[Tensor], Tensor] | None = None, *, layer_class: type[Module] | None = None, output_activation: Callable[[Tensor], Tensor] | None = None, **layer_kwargs: Any)

MLP with standard or custom hidden layers and a linear output layer.

A standard hidden layer is a linear projection followed by activation; a custom one is an instance of layer_class. Subclasses apply custom initialization after super().__init__.

Parameters:

  • in_features (int) –

    Input width.

  • out_features (int) –

    Output width.

  • hidden_features (int | None, default: None ) –

    Hidden width; None uses in_features.

  • hidden_layers (int, default: 2 ) –

    Number of hidden layers.

  • activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied after each standard hidden layer; None uses ReLU. A module instance is shared across the layers.

  • layer_class (type[Module] | None, default: None ) –

    Class of the custom hidden layers, built as layer_class(in_width, out_width, **layer_kwargs). No other activation is applied.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

  • **layer_kwargs (Any, default: {} ) –

    Constructor arguments of the hidden layers only.

Raises:

  • ValueError –

    If hidden_layers < 1, or if activation and layer_class are both given.

Shape
  • Input: (..., in_features).
  • Output: (..., out_features).

Examples:

>>> from inro import SineLayer
>>> model = MLP(2, 3, hidden_features=64, layer_class=SineLayer, omega=30.0)
>>> model(torch.rand(8, 2)).shape
torch.Size([8, 3])
Source code in inro/models/mlp.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    hidden_features: int | None = None,
    hidden_layers: int = 2,
    activation: Callable[[Tensor], Tensor] | None = None,
    *,
    layer_class: type[nn.Module] | None = None,
    output_activation: Callable[[Tensor], Tensor] | None = None,
    **layer_kwargs: Any,
) -> None:
    super().__init__()

    if hidden_layers < 1:
        raise ValueError(f"hidden_layers must be >= 1, got {hidden_layers}")
    if layer_class is None:
        layer_class = nn.Linear
        activation = torch.relu if activation is None else activation
    elif activation is not None:
        raise ValueError("Custom layer_class supplies its own activation.")

    if hidden_features is None:
        hidden_features = in_features

    self.in_features = in_features
    self.out_features = out_features
    self.hidden_features = hidden_features
    self.activation = activation if activation is not None else nn.Identity()

    layers = [layer_class(in_features, hidden_features, **layer_kwargs)]
    for _ in range(hidden_layers - 1):
        layers.append(layer_class(hidden_features, hidden_features, **layer_kwargs))
    layers.append(nn.Linear(hidden_features, out_features))
    self.layers = nn.ModuleList(layers)

    self.output_activation = (
        output_activation if output_activation is not None else nn.Identity()
    )

Random Fourier features

rff

RFF: ReLU MLP with a random Fourier feature input encoding.

RFFEncoding

RFFEncoding(in_features: int, num_frequencies: int = 256, sigma: float = 10.0)

Random Fourier feature encoding of RFF.

Computes \(\sin(2\pi \boldsymbol{B}\boldsymbol{x})\), then \(\cos(2\pi \boldsymbol{B}\boldsymbol{x})\). The fixed buffer B of shape (num_frequencies, C) has its entries drawn from \(\mathcal{N}(0, \sigma^2)\) at construction.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • num_frequencies (int, default: 256 ) –

    Number of random frequency vectors.

  • sigma (float, default: 10.0 ) –

    Standard deviation of the frequencies.

Shape
  • Input: (..., C).
  • Output: (..., 2 * num_frequencies).
Source code in inro/models/rff.py
def __init__(
    self, in_features: int, num_frequencies: int = 256, sigma: float = 10.0
) -> None:
    super().__init__()

    frequencies = torch.randn(num_frequencies, in_features) * sigma
    self.register_buffer("B", frequencies)

    self.out_features = 2 * num_frequencies

RFF

RFF(in_features: int, out_features: int, hidden_features: int, hidden_layers: int = 3, num_frequencies: int = 256, sigma: float = 10.0, output_activation: Callable[[Tensor], Tensor] | None = None)

ReLU MLP with random Fourier features.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • hidden_features (int) –

    Hidden width.

  • hidden_layers (int, default: 3 ) –

    Number of hidden layers.

  • num_frequencies (int, default: 256 ) –

    Number of random frequency vectors.

  • sigma (float, default: 10.0 ) –

    Standard deviation of the frequencies.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
References

Tancik et al., "Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains", NeurIPS 2020.

Source code in inro/models/rff.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    hidden_features: int,
    hidden_layers: int = 3,
    num_frequencies: int = 256,
    sigma: float = 10.0,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    super().__init__()

    self.encoding = RFFEncoding(in_features, num_frequencies, sigma)
    self.mlp = MLP(
        self.encoding.out_features,
        out_features,
        hidden_features,
        hidden_layers,
        layer_class=ReLULayer,
        output_activation=output_activation,
    )

Positional encoding

pemlp

PE-MLP: ReLU MLP with a NeRF-style positional encoding.

PositionalEncoding

PositionalEncoding(in_features: int, num_frequencies: int = 10)

Positional encoding [x, sin(2**k * x), cos(2**k * x)] of PEMLP.

The bands are k = 0, ..., num_frequencies - 1, without the \(\pi\) of NeRF's Eq. (4), and the buffer freqs holds 2**k. The input comes first, then all sines, then all cosines; within each block, bands vary faster than coordinates.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • num_frequencies (int, default: 10 ) –

    Number of frequency bands.

Shape
  • Input: (..., C).
  • Output: (..., C * (2 * num_frequencies + 1)).
Source code in inro/models/pemlp.py
def __init__(self, in_features: int, num_frequencies: int = 10) -> None:
    super().__init__()

    freqs = 2.0 ** torch.linspace(0, num_frequencies - 1, num_frequencies)
    self.register_buffer("freqs", freqs)

    self.out_features = in_features * (2 * num_frequencies + 1)

PEMLP

PEMLP(in_features: int, out_features: int, hidden_features: int, hidden_layers: int = 3, num_frequencies: int = 10, output_activation: Callable[[Tensor], Tensor] | None = None)

ReLU MLP with a NeRF-style positional encoding.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • hidden_features (int) –

    Hidden width.

  • hidden_layers (int, default: 3 ) –

    Number of hidden layers.

  • num_frequencies (int, default: 10 ) –

    Number of frequency bands.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
References

Mildenhall et al., "NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis", ECCV 2020.

Source code in inro/models/pemlp.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    hidden_features: int,
    hidden_layers: int = 3,
    num_frequencies: int = 10,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    super().__init__()

    self.encoding = PositionalEncoding(in_features, num_frequencies)
    self.mlp = MLP(
        self.encoding.out_features,
        out_features,
        hidden_features,
        hidden_layers,
        layer_class=ReLULayer,
        output_activation=output_activation,
    )

SIREN

siren

SIREN: coordinate MLP with sine activations.

SineLayer

SineLayer(in_features: int, out_features: int, omega: float = 30.0)

Sine layer sin(omega * linear(x)) of SIREN.

Parameters:

  • in_features (int) –

    Input width.

  • out_features (int) –

    Output width.

  • omega (float, default: 30.0 ) –

    Frequency multiplier.

Shape
  • Input: (..., in_features).
  • Output: (..., out_features).
Source code in inro/models/siren.py
def __init__(
    self, in_features: int, out_features: int, omega: float = 30.0
) -> None:
    super().__init__()

    self.in_features = in_features
    self.omega = omega

    self.linear = nn.Linear(in_features, out_features)

SIREN

SIREN(in_features: int, out_features: int, hidden_features: int, hidden_layers: int = 3, omega: float = 30.0, first_omega: float | None = None, output_activation: Callable[[Tensor], Tensor] | None = None)

Sinusoidal representation network.

An MLP with SineLayer hidden layers, a linear output layer and SIREN weight initialization. Biases keep PyTorch's defaults.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • hidden_features (int) –

    Hidden width.

  • hidden_layers (int, default: 3 ) –

    Number of hidden layers.

  • omega (float, default: 30.0 ) –

    Frequency multiplier.

  • first_omega (float | None, default: None ) –

    Frequency multiplier of the first layer alone, the one inro.fresh_select chooses; None uses omega. The later layers' initialization follows omega.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
References

Sitzmann et al., "Implicit Neural Representations with Periodic Activation Functions", NeurIPS 2020.

Source code in inro/models/siren.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    hidden_features: int,
    hidden_layers: int = 3,
    omega: float = 30.0,
    first_omega: float | None = None,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    super().__init__(
        in_features,
        out_features,
        hidden_features,
        hidden_layers,
        layer_class=SineLayer,
        omega=omega,
        output_activation=output_activation,
    )
    self.omega = omega
    if first_omega is not None:
        self.layers[0].omega = first_omega
    self.reset_parameters()

reset_parameters

reset_parameters() -> None

Initialize weights with SIREN bounds, leaving biases unchanged.

Source code in inro/models/siren.py
def reset_parameters(self) -> None:
    """Initialize weights with SIREN bounds, leaving biases unchanged."""
    _siren_init_(self.layers, self.omega)

FreSh

fresh

FreSh: a SIREN's first-layer frequency chosen from the signal's spectrum.

References

Kania et al., "FreSh: Frequency Shifting for Accelerated Neural Representation Learning", ICLR 2025.

fresh_spectrum

fresh_spectrum(signal: Tensor, size: int = 64) -> Tensor

Compute the spectrum FreSh compares, as a distribution.

As in the authors' code, the grid is first resized to a square or a cube of its smallest side, at most 1024: bilinearly with antialiasing in 2D, trilinearly in 3D. The amplitudes of its discrete Fourier transform are summed over the channels, and those of nonnegative frequencies along shells of equal index sum d.

Parameters:

  • signal (Tensor) –

    Values on a grid, (H, W, channels) or (depth, H, W, channels).

  • size (int, default: 64 ) –

    Number of shells d = 1, ..., size, which leave out the constant term, at most half the side after resizing.

Returns:

  • Tensor –

    The shells divided by their sum, a float32 tensor (size,) that sums to one.

Raises:

  • ValueError –

    If the grid is neither 2D nor 3D, or size is too large.

Source code in inro/models/fresh.py
def fresh_spectrum(signal: Tensor, size: int = 64) -> Tensor:
    """Compute the spectrum FreSh compares, as a distribution.

    As in the authors' code, the grid is first resized to a square or a cube of its
    smallest side, at most ``1024``: bilinearly with antialiasing in 2D, trilinearly in
    3D. The amplitudes of its discrete Fourier transform are summed over the channels,
    and those of nonnegative frequencies along shells of equal index sum ``d``.

    Args:
        signal: Values on a grid, ``(H, W, channels)`` or ``(depth, H, W, channels)``.
        size: Number of shells ``d = 1, ..., size``, which leave out the constant term,
            at most half the side after resizing.

    Returns:
        The shells divided by their sum, a float32 tensor ``(size,)`` that sums to one.

    Raises:
        ValueError: If the grid is neither 2D nor 3D, or ``size`` is too large.
    """
    grid = signal.detach().movedim(-1, 0).float()  # (channels, *sides)
    sides = grid.shape[1:]
    if len(sides) not in (2, 3):
        raise ValueError(f"expected a 2D or 3D grid, got {len(sides)} sides")
    side = min(min(sides), 1024)
    if any(length != side for length in sides):
        planar = len(sides) == 2
        grid = F.interpolate(
            grid[None],
            size=[side] * len(sides),
            mode="bilinear" if planar else "trilinear",
            align_corners=False,
            antialias=planar,
        )[0]
    if size > side // 2:
        raise ValueError(f"{size} shells reach past half the side, {side}")
    axes = tuple(range(1, grid.ndim))
    amplitude = torch.fft.fftn(grid, dim=axes).abs().sum(0)
    corner = amplitude[(slice(0, size + 1),) * len(sides)]
    steps = torch.arange(size + 1, device=corner.device)
    shell = sum(
        steps.view([-1 if axis == i else 1 for axis in range(len(sides))])
        for i in range(len(sides))
    )
    sums = torch.zeros(len(sides) * size + 1, device=corner.device)
    sums.scatter_add_(0, shell.flatten(), corner.flatten())
    spectrum = sums[1 : size + 1]
    return spectrum / spectrum.sum()

fresh_distance

fresh_distance(p: Tensor, q: Tensor) -> float

Return the Wasserstein-1 distance of two spectra.

Parameters:

  • p (Tensor) –

    Spectrum (n,) on the points 0, ..., n - 1.

  • q (Tensor) –

    Spectrum (n,) on the same points.

Source code in inro/models/fresh.py
def fresh_distance(p: Tensor, q: Tensor) -> float:
    """Return the Wasserstein-1 distance of two spectra.

    Args:
        p: Spectrum ``(n,)`` on the points ``0, ..., n - 1``.
        q: Spectrum ``(n,)`` on the same points.
    """
    return float((p.cumsum(0) - q.cumsum(0))[:-1].abs().sum())

fresh_select

fresh_select(targets: Sequence[Tensor], render: Callable[[float], Tensor], omegas: Sequence[float] = tuple(range(10, 201, 10)), size: int = 64, inits: int = 10, seed: int = 0) -> tuple[float, dict[float, float]]

Choose a SIREN's first-layer frequency with FreSh.

A candidate's score is the mean fresh_distance between the fresh_spectrum of its i-th untrained output and that of targets[i % len(targets)].

Parameters:

  • targets (Sequence[Tensor]) –

    Signal on its grid, as fresh_spectrum takes it, or several views or frames.

  • render (Callable[[float], Tensor]) –

    Function that builds a new, untrained model whose first layer has the given frequency and returns its output on the signal's grid.

  • omegas (Sequence[float], default: tuple(range(10, 201, 10)) ) –

    Candidate frequencies.

  • size (int, default: 64 ) –

    Number of spectrum shells compared.

  • inits (int, default: 10 ) –

    Number of initializations averaged per candidate.

  • seed (int, default: 0 ) –

    Seed of the global random state, set anew for every candidate.

Returns:

  • tuple[float, dict[float, float]] –

    The candidate of least score, the smaller one on a tie, and every candidate's score.

Examples:

>>> from inro import SIREN, create_coordinates
>>> x = create_coordinates((32, 32))
>>> omega, scores = fresh_select(
...     [torch.rand(32, 32, 3)],
...     lambda omega: SIREN(2, 3, 16, first_omega=omega)(x),
...     omegas=(10, 20),
...     size=8,
...     inits=2,
... )
>>> omega in scores
True
Source code in inro/models/fresh.py
@torch.no_grad()
def fresh_select(
    targets: Sequence[Tensor],
    render: Callable[[float], Tensor],
    omegas: Sequence[float] = tuple(range(10, 201, 10)),
    size: int = 64,
    inits: int = 10,
    seed: int = 0,
) -> tuple[float, dict[float, float]]:
    """Choose a SIREN's first-layer frequency with FreSh.

    A candidate's score is the mean :func:`fresh_distance` between the
    :func:`fresh_spectrum` of its ``i``-th untrained output and that of
    ``targets[i % len(targets)]``.

    Args:
        targets: Signal on its grid, as :func:`fresh_spectrum` takes it, or several
            views or frames.
        render: Function that builds a new, untrained model whose first layer has the
            given frequency and returns its output on the signal's grid.
        omegas: Candidate frequencies.
        size: Number of spectrum shells compared.
        inits: Number of initializations averaged per candidate.
        seed: Seed of the global random state, set anew for every candidate.

    Returns:
        The candidate of least score, the smaller one on a tie, and every candidate's
        score.

    Examples:
        >>> from inro import SIREN, create_coordinates
        >>> x = create_coordinates((32, 32))
        >>> omega, scores = fresh_select(
        ...     [torch.rand(32, 32, 3)],
        ...     lambda omega: SIREN(2, 3, 16, first_omega=omega)(x),
        ...     omegas=(10, 20),
        ...     size=8,
        ...     inits=2,
        ... )
        >>> omega in scores
        True
    """
    goals = [fresh_spectrum(target, size) for target in targets]
    scores = {}
    for omega in omegas:
        torch.manual_seed(seed)
        distances = [
            fresh_distance(goals[i % len(goals)], fresh_spectrum(render(omega), size))
            for i in range(inits)
        ]
        scores[omega] = sum(distances) / inits
    return min(omegas, key=lambda omega: (scores[omega], omega)), scores

Gauss

gauss

Gauss: coordinate MLP with Gaussian activations.

GaussLayer

GaussLayer(in_features: int, out_features: int, scale: float = 10.0)

Gaussian layer exp(-(scale * linear(x))**2) of Gauss.

Parameters:

  • in_features (int) –

    Input width.

  • out_features (int) –

    Output width.

  • scale (float, default: 10.0 ) –

    Inverse Gaussian width.

Shape
  • Input: (..., in_features).
  • Output: (..., out_features).
Source code in inro/models/gauss.py
def __init__(
    self, in_features: int, out_features: int, scale: float = 10.0
) -> None:
    super().__init__()

    self.in_features = in_features
    self.scale = scale

    self.linear = nn.Linear(in_features, out_features)

Gauss

Gauss(in_features: int, out_features: int, hidden_features: int, hidden_layers: int = 3, scale: float = 10.0, output_activation: Callable[[Tensor], Tensor] | None = None)

Gaussian coordinate MLP.

An MLP with GaussLayer hidden layers, a linear output layer and PyTorch's default initialization.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • hidden_features (int) –

    Hidden width.

  • hidden_layers (int, default: 3 ) –

    Number of hidden layers.

  • scale (float, default: 10.0 ) –

    Inverse Gaussian width.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
References

Ramasinghe and Lucey, "Beyond Periodicity: Towards a Unifying Framework for Activations in Coordinate-MLPs", ECCV 2022.

Source code in inro/models/gauss.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    hidden_features: int,
    hidden_layers: int = 3,
    scale: float = 10.0,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    super().__init__(
        in_features,
        out_features,
        hidden_features,
        hidden_layers,
        layer_class=GaussLayer,
        scale=scale,
        output_activation=output_activation,
    )

WIRE

wire

WIRE: coordinate MLPs with Gabor wavelet activations (real and complex).

RealGaborLayer

RealGaborLayer(in_features: int, out_features: int, omega: float = 20.0, scale: float = 10.0)

Real Gabor layer cos(omega * z1) * exp(-(scale * z2)**2).

A linear projection gives z1 and z2, with out_features channels each.

Parameters:

  • in_features (int) –

    Input width.

  • out_features (int) –

    Output width.

  • omega (float, default: 20.0 ) –

    Frequency multiplier.

  • scale (float, default: 10.0 ) –

    Inverse Gaussian width.

Shape
  • Input: (..., in_features).
  • Output: (..., out_features).
Source code in inro/models/wire.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    omega: float = 20.0,
    scale: float = 10.0,
) -> None:
    super().__init__()
    self.in_features = in_features
    self.omega = omega
    self.scale = scale

    self.linear = nn.Linear(in_features, 2 * out_features)

ComplexGaborLayer

ComplexGaborLayer(in_features: int, out_features: int, omega: float = 20.0, scale: float = 10.0, is_first: bool = False)

Complex Gabor layer exp(1j * omega * z - abs(scale * z)**2).

Here z = linear(x).

Parameters:

  • in_features (int) –

    Input width.

  • out_features (int) –

    Output width.

  • omega (float, default: 20.0 ) –

    Frequency multiplier.

  • scale (float, default: 10.0 ) –

    Inverse Gaussian width.

  • is_first (bool, default: False ) –

    Whether to use real weights, for real inputs, rather than complex.

Shape
  • Input: (..., in_features), real if is_first, else complex.
  • Output: (..., out_features), complex.
Source code in inro/models/wire.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    omega: float = 20.0,
    scale: float = 10.0,
    is_first: bool = False,
) -> None:
    super().__init__()
    self.in_features = in_features
    self.omega = omega
    self.scale = scale

    dtype = torch.float if is_first else torch.cfloat
    self.linear = nn.Linear(in_features, out_features, dtype=dtype)

RealWIRE

RealWIRE(in_features: int, out_features: int, hidden_features: int, hidden_layers: int = 3, omega: float = 20.0, scale: float = 10.0, output_activation: Callable[[Tensor], Tensor] | None = None)

Real-valued variant of WIRE with RealGaborLayer hidden layers.

Uses a linear output layer and PyTorch's default initialization. Each hidden layer projects to twice the reduced width, to form its two wavelet terms.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • hidden_features (int) –

    Hidden width, divided by sqrt(2) and rounded down.

  • hidden_layers (int, default: 3 ) –

    Number of hidden layers.

  • omega (float, default: 20.0 ) –

    Frequency multiplier.

  • scale (float, default: 10.0 ) –

    Inverse Gaussian width.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
Source code in inro/models/wire.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    hidden_features: int,
    hidden_layers: int = 3,
    omega: float = 20.0,
    scale: float = 10.0,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    hidden_features = int(hidden_features / math.sqrt(2))

    super().__init__(
        in_features,
        out_features,
        hidden_features,
        hidden_layers,
        layer_class=RealGaborLayer,
        omega=omega,
        scale=scale,
        output_activation=output_activation,
    )

WIRE

WIRE(in_features: int, out_features: int, hidden_features: int, hidden_layers: int = 3, omega: float = 20.0, scale: float = 10.0, output_activation: Callable[[Tensor], Tensor] | None = None)

Wavelet implicit neural representation with complex Gabor layers.

Uses ComplexGaborLayer hidden layers, the first with real weights, and a complex linear output layer, all with PyTorch's default initialization. The real part of the output goes through output_activation. Dividing the width keeps the parameter count comparable to real-valued networks.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • hidden_features (int) –

    Hidden width, divided by sqrt(2) and rounded down.

  • hidden_layers (int, default: 3 ) –

    Number of hidden layers.

  • omega (float, default: 20.0 ) –

    Frequency multiplier.

  • scale (float, default: 10.0 ) –

    Inverse Gaussian width.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1], real.
  • Output: (..., D), real.
References

Saragadam et al., "WIRE: Wavelet Implicit Neural Representations", CVPR 2023.

Source code in inro/models/wire.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    hidden_features: int,
    hidden_layers: int = 3,
    omega: float = 20.0,
    scale: float = 10.0,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    super().__init__()
    if hidden_layers < 1:
        raise ValueError(f"hidden_layers must be >= 1, got {hidden_layers}")

    # Account for the two real components of each complex parameter.
    hidden_features = int(hidden_features / math.sqrt(2))

    self.layers = nn.ModuleList()
    self.layers.append(
        ComplexGaborLayer(
            in_features, hidden_features, omega=omega, scale=scale, is_first=True
        )
    )
    for _ in range(hidden_layers - 1):
        self.layers.append(
            ComplexGaborLayer(
                hidden_features, hidden_features, omega=omega, scale=scale
            )
        )
    self.layers.append(nn.Linear(hidden_features, out_features, dtype=torch.cfloat))

    self.output_activation = (
        output_activation if output_activation is not None else nn.Identity()
    )

FINER

finer

FINER: coordinate MLP with variable-periodic sine activations.

FinerLayer

FinerLayer(in_features: int, out_features: int, omega: float = 30.0)

Variable-periodic sine layer sin(omega * (abs(z) + 1) * z) of FINER.

Here z = linear(x). The scale abs(z) + 1 is detached, which keeps its derivative out of the gradient.

Parameters:

  • in_features (int) –

    Input width.

  • out_features (int) –

    Output width.

  • omega (float, default: 30.0 ) –

    Frequency multiplier.

Shape
  • Input: (..., in_features).
  • Output: (..., out_features).
Source code in inro/models/finer.py
def __init__(
    self, in_features: int, out_features: int, omega: float = 30.0
) -> None:
    super().__init__()

    self.in_features = in_features
    self.omega = omega

    self.linear = nn.Linear(in_features, out_features)

FINER

FINER(in_features: int, out_features: int, hidden_features: int, hidden_layers: int = 3, omega: float = 30.0, output_activation: Callable[[Tensor], Tensor] | None = None)

Variable-periodic coordinate MLP.

An MLP with FinerLayer hidden layers, a linear output layer and SIREN weight initialization. Biases keep PyTorch's defaults; the paper's enlarged first-layer bias range is not exposed.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • hidden_features (int) –

    Hidden width.

  • hidden_layers (int, default: 3 ) –

    Number of hidden layers.

  • omega (float, default: 30.0 ) –

    Frequency multiplier.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
References

Liu et al., "FINER: Flexible Spectral-bias Tuning in Implicit NEural Representation by Variable-periodic Activation Functions", CVPR 2024.

Source code in inro/models/finer.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    hidden_features: int,
    hidden_layers: int = 3,
    omega: float = 30.0,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    super().__init__(
        in_features,
        out_features,
        hidden_features,
        hidden_layers,
        layer_class=FinerLayer,
        omega=omega,
        output_activation=output_activation,
    )
    self.reset_parameters()

reset_parameters

reset_parameters() -> None

Initialize weights with SIREN bounds, leaving biases unchanged.

Source code in inro/models/finer.py
def reset_parameters(self) -> None:
    """Initialize weights with SIREN bounds, leaving biases unchanged."""
    _siren_init_(self.layers, self.layers[0].omega)

Multiplicative filter networks

mfn

MFN: multiplicative filter network with Gabor filters.

GaborFilter

GaborFilter(in_features: int, out_features: int, input_scale: float, alpha: float = 1.0, beta: float = 1.0)

Gabor filter sin(Wx + b) * exp(-gamma * ||x - mu||**2 / 2) of MFN.

The centers mu start uniformly in [-1, 1] and the inverse squared widths gamma from a Gamma distribution, both learnable. The weights W are scaled by input_scale * sqrt(gamma) and the biases b start in \([-\pi, \pi]\).

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of filters.

  • input_scale (float) –

    Frequency scale.

  • alpha (float, default: 1.0 ) –

    Gamma shape.

  • beta (float, default: 1.0 ) –

    Gamma rate.

Shape
  • Input: (..., C), with at least one leading dimension.
  • Output: (..., out_features).
Source code in inro/models/mfn.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    input_scale: float,
    alpha: float = 1.0,
    beta: float = 1.0,
) -> None:
    super().__init__()

    self.linear = nn.Linear(in_features, out_features)
    self.mu = nn.Parameter(2 * torch.rand(out_features, in_features) - 1)
    self.gamma = nn.Parameter(
        torch.distributions.Gamma(alpha, beta).sample((out_features,))
    )
    with torch.no_grad():
        self.linear.weight *= input_scale * self.gamma[:, None].sqrt()
        self.linear.bias.uniform_(-math.pi, math.pi)

MFN

MFN(in_features: int, out_features: int, hidden_features: int, hidden_layers: int = 4, input_scale: float = 256.0, weight_scale: float = 1.0, alpha: float = 6.0, beta: float = 1.0, output_activation: Callable[[Tensor], Tensor] | None = None)

Multiplicative Gabor filter network.

Each hidden layer multiplies a GaborFilter of the raw coordinates by a linear transform of the previous hidden state; the first takes its filter alone. A linear layer gives the output.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • hidden_features (int) –

    Hidden width.

  • hidden_layers (int, default: 4 ) –

    Number of hidden layers, one filter each.

  • input_scale (float, default: 256.0 ) –

    Frequency scale, divided by sqrt(hidden_layers) per filter.

  • weight_scale (float, default: 1.0 ) –

    Scale of the hidden linear weights, drawn uniformly within sqrt(weight_scale / hidden_features) of zero.

  • alpha (float, default: 6.0 ) –

    Gamma shape, divided by hidden_layers per filter.

  • beta (float, default: 1.0 ) –

    Gamma rate.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1], with at least one leading dimension.
  • Output: (..., D).
References

Fathony et al., "Multiplicative Filter Networks", ICLR 2021.

Source code in inro/models/mfn.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    hidden_features: int,
    hidden_layers: int = 4,
    input_scale: float = 256.0,
    weight_scale: float = 1.0,
    alpha: float = 6.0,
    beta: float = 1.0,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    super().__init__()
    if hidden_layers < 1:
        raise ValueError(f"hidden_layers must be >= 1, got {hidden_layers}")

    self.filters = nn.ModuleList(
        [
            GaborFilter(
                in_features,
                hidden_features,
                input_scale / math.sqrt(hidden_layers),
                alpha / hidden_layers,
                beta,
            )
            for _ in range(hidden_layers)
        ]
    )
    self.linears = nn.ModuleList(
        [
            nn.Linear(hidden_features, hidden_features)
            for _ in range(hidden_layers - 1)
        ]
    )
    self.output_linear = nn.Linear(hidden_features, out_features)
    self.output_activation = (
        output_activation if output_activation is not None else nn.Identity()
    )

    bound = math.sqrt(weight_scale / hidden_features)
    with torch.no_grad():
        for linear in self.linears:
            linear.weight.uniform_(-bound, bound)

Instant-NGP

instant_ngp

Instant-NGP: a multiresolution hash encoding and an MLP, in pure PyTorch.

The encoding reproduces tiny-cuda-nn's HashGrid: level scales, half-cell offsets, table sizes and the hash function match the reference, so equal table contents give equal features. Only the storage layout of the concatenated table differs.

HashEncoding

HashEncoding(in_features: int, num_levels: int = 16, features_per_level: int = 2, log2_hashmap_size: int = 19, base_resolution: int = 16, max_resolution: int = 2048)

Multiresolution hash encoding of InstantNGP.

Level l scales unit-cube positions by N_min * b**l - 1, with growth b = (N_max / N_min) ** (1 / (L - 1)), and adds a half-cell offset that staggers the levels (Appendix A of the paper). Its grid has ceil(scale) + 1 vertices per axis, indexed densely when they fit in the table and hashed otherwise. The features of the 2**C surrounding vertices are interpolated multilinearly and concatenated across levels.

Parameters:

  • in_features (int) –

    Number of coordinate axes C, from 1 to 3.

  • num_levels (int, default: 16 ) –

    Number of resolution levels L.

  • features_per_level (int, default: 2 ) –

    Feature channels F per table entry.

  • log2_hashmap_size (int, default: 19 ) –

    Base-2 logarithm of the maximum table entries per level.

  • base_resolution (int, default: 16 ) –

    Coarsest resolution N_min.

  • max_resolution (int, default: 2048 ) –

    Finest resolution N_max over [-1, 1]; the reference uses 2048 for NeRF and SDFs, max(H, W) / 2 for images and the voxel resolution for volumes.

Attributes:

  • embeddings –

    Learnable tables of all levels, concatenated, that start uniformly in [-1e-4, 1e-4].

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., L * F).
Source code in inro/models/instant_ngp.py
def __init__(
    self,
    in_features: int,
    num_levels: int = 16,
    features_per_level: int = 2,
    log2_hashmap_size: int = 19,
    base_resolution: int = 16,
    max_resolution: int = 2048,
) -> None:
    super().__init__()
    if not 1 <= in_features <= len(_PRIMES):
        raise ValueError(f"in_features must be in [1, {len(_PRIMES)}]")
    self.in_features = in_features
    self.num_levels = num_levels
    self.features_per_level = features_per_level

    self.scales = _level_scales(num_levels, base_resolution, max_resolution)
    self.resolutions = [math.ceil(scale) + 1 for scale in self.scales]
    self.table_sizes = [
        -(-min(resolution**in_features, 2**log2_hashmap_size) // 8) * 8
        for resolution in self.resolutions
    ]
    self.offsets = [0]
    for size in self.table_sizes[:-1]:
        self.offsets.append(self.offsets[-1] + size)
    # The per-level constants as tensors, for the batched forward pass. A level
    # whose vertices outnumber its table is hashed, the others are indexed densely.
    levels = {
        "level_scales": torch.tensor(self.scales, dtype=torch.float32),
        "level_resolutions": torch.tensor(self.resolutions, dtype=torch.long),
        "level_table_sizes": torch.tensor(self.table_sizes, dtype=torch.long),
        "level_offsets": torch.tensor(self.offsets, dtype=torch.long),
        "level_hashed": torch.tensor(
            [
                size < resolution**in_features
                for resolution, size in zip(self.resolutions, self.table_sizes)
            ]
        ),
    }
    for name, value in levels.items():
        self.register_buffer(name, value, persistent=False)

    self.register_buffer(
        "corners",
        torch.tensor(list(product((0, 1), repeat=in_features)), dtype=torch.long),
        persistent=False,
    )
    self.register_buffer(
        "primes",
        torch.tensor(_PRIMES[:in_features], dtype=torch.long),
        persistent=False,
    )
    self.embeddings = nn.Parameter(
        torch.empty(sum(self.table_sizes), features_per_level).uniform_(-1e-4, 1e-4)
    )

out_features property

out_features: int

Number of output features, num_levels * features_per_level.

InstantNGP

InstantNGP(in_features: int, out_features: int, num_levels: int = 16, features_per_level: int = 2, log2_hashmap_size: int = 19, base_resolution: int = 16, max_resolution: int = 2048, hidden_features: int = 64, hidden_layers: int = 1, output_activation: Callable[[Tensor], Tensor] | None = None)

Instant neural graphics primitive: a hash encoding and a bias-free ReLU MLP.

Mirrors tiny-cuda-nn's HashGrid followed by FullyFusedMLP, which has no biases. The reference uses one 64-wide hidden layer for the NeRF density network and two for images and SDFs.

Parameters:

  • in_features (int) –

    Number of coordinate axes C, from 1 to 3.

  • out_features (int) –

    Number of output channels D.

  • num_levels (int, default: 16 ) –

    Number of resolution levels L.

  • features_per_level (int, default: 2 ) –

    Feature channels F per table entry.

  • log2_hashmap_size (int, default: 19 ) –

    Base-2 logarithm of the maximum table entries per level.

  • base_resolution (int, default: 16 ) –

    Coarsest resolution N_min.

  • max_resolution (int, default: 2048 ) –

    Finest resolution N_max over [-1, 1]; the reference uses 2048 for NeRF and SDFs, max(H, W) / 2 for images and the voxel resolution for volumes.

  • hidden_features (int, default: 64 ) –

    Hidden width.

  • hidden_layers (int, default: 1 ) –

    Number of hidden layers.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
Note

The reference trains with Adam, betas=(0.9, 0.99) and eps=1e-15; the tiny epsilon matters for rarely updated table entries.

References

Müller et al., "Instant Neural Graphics Primitives with a Multiresolution Hash Encoding", SIGGRAPH 2022.

Source code in inro/models/instant_ngp.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    num_levels: int = 16,
    features_per_level: int = 2,
    log2_hashmap_size: int = 19,
    base_resolution: int = 16,
    max_resolution: int = 2048,
    hidden_features: int = 64,
    hidden_layers: int = 1,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    super().__init__()
    self.in_features = in_features
    self.out_features = out_features
    self.output_activation = (
        output_activation if output_activation is not None else nn.Identity()
    )

    self.encoding = HashEncoding(
        in_features,
        num_levels=num_levels,
        features_per_level=features_per_level,
        log2_hashmap_size=log2_hashmap_size,
        base_resolution=base_resolution,
        max_resolution=max_resolution,
    )
    self.decoder = MLP(
        self.encoding.out_features,
        out_features,
        hidden_features,
        hidden_layers,
        bias=False,  # hidden layers
    )
    # FullyFusedMLP has no output bias either.
    output_layer = self.decoder.layers[-1]
    self.decoder.layers[-1] = nn.Linear(
        output_layer.in_features, out_features, bias=False
    )

Feature grids

multivector

Learnable feature grids and expressions that combine lines, planes and volumes.

FeatureGrid

FeatureGrid(features: int, resolution: int | Sequence[int])

Learnable feature grid with linear, bilinear or trilinear interpolation.

Coordinates span [-1, 1] with align_corners=True, and points outside read the border.

Parameters:

  • features (int) –

    Feature channels at each grid vertex.

  • resolution (int | Sequence[int]) –

    Resolution per axis, an int for a 1D grid or a sequence of up to three. The order (H, W) or (depth, H, W) is reversed against the coordinates: their first column indexes W, the second H and the third depth.

Attributes:

  • grid –

    Learnable features (1, features, *resolution), a 1D grid with a dummy height axis for grid_sample.

Shape
  • Input: (N, C) coordinates in [-1, 1], or (N,) in 1D.
  • Output: (N, features).

Examples:

>>> grid = FeatureGrid(features=8, resolution=(32, 64))
>>> grid(torch.rand(100, 2) * 2 - 1).shape
torch.Size([100, 8])
Source code in inro/models/multivector.py
def __init__(self, features: int, resolution: int | Sequence[int]) -> None:
    super().__init__()
    self.features = features

    if isinstance(resolution, int):
        resolution = (resolution,)
    resolution = tuple(resolution)
    self.in_features = len(resolution)

    if self.in_features == 1:
        # grid_sample requires at least two spatial dimensions.
        resolution = (1,) + resolution
    elif self.in_features not in (2, 3):
        raise ValueError(f"FeatureGrid supports 1D–3D, got {self.in_features}D")

    self.grid = nn.Parameter(torch.empty(1, features, *resolution))
    self.reset_parameters()

reset_parameters

reset_parameters() -> None

Re-initialize the grid features from \(\mathcal{N}(0, 0.1^2)\).

Source code in inro/models/multivector.py
def reset_parameters(self) -> None:
    r"""Re-initialize the grid features from :math:`\mathcal{N}(0, 0.1^2)`."""
    nn.init.normal_(self.grid, std=0.1)

MultiVector

MultiVector(in_features: int, expr: str, features: int | Sequence[int], resolution: int | Sequence[int])

Learnable line, plane and volume grids combined by an expression.

Each atom samples a FeatureGrid at its coordinate columns, following the grid decomposition of GA-Planes:

  • e1, e2, e3: lines along the first, second and third columns.
  • e12, e13, e23: planes over the named pair of columns.
  • e123: a volume over all three columns.
  • o: elementwise product; +: elementwise sum.
  • [term, term, ...]: channel concatenation.

Precedence is o > + > ,. Atoms on the third axis need in_features=3; neither this nor the operands' widths are checked at construction.

Parameters:

  • in_features (int) –

    Number of coordinate axes C, 2 or 3.

  • expr (str) –

    Expression that selects and combines the grids.

  • features (int | Sequence[int]) –

    Feature channels per grid: an int for all grids, or (line, plane, volume) values, a short sequence repeating its last value. Operands of o and + must have equal channel counts.

  • resolution (int | Sequence[int]) –

    Resolution per axis per grid, given as features is. Planes are square and volumes are cubic.

Attributes:

  • lines –

    Line grids keyed by atom name, or None when unused.

  • planes –

    Plane grids keyed by atom name, or None when unused.

  • volume –

    Volume grid, or None when unused.

  • features_per_term –

    Width of each top-level concatenated term, or one width when the expression is not bracketed.

Shape
  • Input: (N, C) coordinates in [-1, 1].
  • Output: (N, out_features).

Examples:

>>> model = MultiVector(2, "[e1 o e2, e12]", 8, (64, 16))
>>> model(torch.rand(100, 2) * 2 - 1).shape
torch.Size([100, 16])
References

Sivgin et al., "Geometric Algebra Planes: Convex Implicit Neural Volumes", ICML 2025.

Source code in inro/models/multivector.py
def __init__(
    self,
    in_features: int,
    expr: str,
    features: int | Sequence[int],
    resolution: int | Sequence[int],
) -> None:
    super().__init__()

    if in_features not in (2, 3):
        raise ValueError(f"in_features must be 2 or 3, got {in_features}")
    self.in_features = in_features
    features = _per_grid(features)
    resolution = _per_grid(resolution)

    self._ast = _parse_expression(expr)

    atoms = _collect_atoms(self._ast)
    parsed_atoms = {name: _parse_atom(name) for name in atoms}
    line_axes = sorted(
        {axes[0] for kind, axes in parsed_atoms.values() if kind == "line"}
    )
    plane_pairs = sorted(
        {axes for kind, axes in parsed_atoms.values() if kind == "plane"}
    )
    needs_volume = any(kind == "volume" for kind, _ in parsed_atoms.values())

    self.lines = (
        nn.ModuleDict(
            {
                f"e{axis + 1}": FeatureGrid(features[0], resolution[0])
                for axis in line_axes
            }
        )
        if line_axes
        else None
    )
    self.planes = (
        nn.ModuleDict(
            {
                f"e{first + 1}{second + 1}": FeatureGrid(
                    features[1], (resolution[1],) * 2
                )
                for first, second in plane_pairs
            }
        )
        if plane_pairs
        else None
    )
    self.volume = (
        FeatureGrid(features[2], (resolution[2],) * 3) if needs_volume else None
    )

    self._line_axes = {f"e{axis + 1}": axis for axis in line_axes}
    self._plane_pairs = {
        f"e{first + 1}{second + 1}": [first, second]
        for first, second in plane_pairs
    }

    grid = {"line": 0, "plane": 1, "volume": 2}
    self._atom_features = {
        name: features[grid[kind]] for name, (kind, _) in parsed_atoms.items()
    }

    # A redundant draw, kept so that seeded models reproduce their checkpoints.
    self.reset_parameters()

    self._terms = self._ast[1] if self._ast[0] == "concat" else [self._ast]
    self.features_per_term = [
        _count_features(term, self._atom_features) for term in self._terms
    ]

out_features property

out_features: int

Number of output features, sum(features_per_term).

reset_parameters

reset_parameters() -> None

Re-initialize every grid's features from \(\mathcal{N}(0, 0.1^2)\).

Source code in inro/models/multivector.py
def reset_parameters(self) -> None:
    r"""Re-initialize every grid's features from :math:`\mathcal{N}(0, 0.1^2)`."""
    for grid in self.modules():
        if isinstance(grid, FeatureGrid):
            grid.reset_parameters()

GA-Planes

gaplanes

GA-Planes: geometric algebra planes, feature grids with an MLP decoder.

GAPlanes

GAPlanes(in_features: int, out_features: int, expr: str | None = None, features: int | Sequence[int] = 32, resolution: int | Sequence[int] = (256, 64, 16), hidden_features: int = 64, hidden_layers: int = 2, output_activation: Callable[[Tensor], Tensor] | None = None)

Line, plane and volume feature grids with a ReLU MLP decoder.

The default MultiVector expression combines a low-rank term with lower-resolution grids: in 2D it concatenates the product of two line grids and a plane; in 3D it also includes line-plane products and a volume.

Parameters:

  • in_features (int) –

    Number of coordinate axes C, 2 or 3.

  • out_features (int) –

    Number of output channels D.

  • expr (str | None, default: None ) –

    MultiVector expression; None selects all grids.

  • features (int | Sequence[int], default: 32 ) –

    Feature channels per grid: an int for all grids, or (line, plane, volume) values, a short sequence repeating its last value. Operands of o and + must have equal channel counts.

  • resolution (int | Sequence[int], default: (256, 64, 16) ) –

    Resolution per axis per grid, given as features is. Planes are square and volumes are cubic.

  • hidden_features (int, default: 64 ) –

    Hidden width.

  • hidden_layers (int, default: 2 ) –

    Number of hidden layers.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
References

Sivgin et al., "Geometric Algebra Planes: Convex Implicit Neural Volumes", ICML 2025.

Source code in inro/models/gaplanes.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    expr: str | None = None,
    features: int | Sequence[int] = 32,
    resolution: int | Sequence[int] = (256, 64, 16),
    hidden_features: int = 64,
    hidden_layers: int = 2,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    super().__init__()
    if in_features not in (2, 3):
        raise ValueError(f"in_features must be 2 or 3, got {in_features}")
    self.in_features = in_features
    self.out_features = out_features
    self.output_activation = (
        output_activation if output_activation is not None else nn.Identity()
    )

    self.mv = MultiVector(
        in_features=in_features,
        expr=self._DEFAULT_EXPR[in_features] if expr is None else expr,
        features=features,
        resolution=resolution,
    )
    self.decoder = MLP(
        self.mv.out_features, out_features, hidden_features, hidden_layers
    )

TensoRF

tensorf

TensoRF: tensor factorizations expressed as GA-Planes grid encodings.

TensoRF

TensoRF(in_features: int, out_features: int, rank: int = 16, resolution: int = 128, mode: Literal['cp', 'vm'] = 'vm', hidden_features: int = 128, hidden_layers: int = 2, output_activation: Callable[[Tensor], Tensor] | None = None)

Tensor-factorized grid INR with a ReLU MLP decoder.

A GAPlanes model with one resolution and one feature count for all grids. CP multiplies line grids into rank features; VM concatenates three line-plane products into 3 * rank features.

Parameters:

  • in_features (int) –

    Number of coordinate axes C, 2 or 3.

  • out_features (int) –

    Number of output channels D.

  • rank (int, default: 16 ) –

    Feature channels per grid.

  • resolution (int, default: 128 ) –

    Resolution per axis of each grid.

  • mode (Literal['cp', 'vm'], default: 'vm' ) –

    Factorization, "cp" or "vm"; "vm" becomes "cp" in 2D.

  • hidden_features (int, default: 128 ) –

    Hidden width.

  • hidden_layers (int, default: 2 ) –

    Number of hidden layers.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
References

Chen et al., "TensoRF: Tensorial Radiance Fields", ECCV 2022.

Source code in inro/models/tensorf.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    rank: int = 16,
    resolution: int = 128,
    mode: Literal["cp", "vm"] = "vm",
    hidden_features: int = 128,
    hidden_layers: int = 2,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    if mode not in ("cp", "vm"):
        raise ValueError(f"mode must be 'cp' or 'vm', got {mode!r}")
    if in_features == 2:
        mode = "cp"  # VM needs a third axis

    super().__init__(
        in_features=in_features,
        out_features=out_features,
        expr=self._MODE_EXPR.get((in_features, mode)),
        features=rank,
        resolution=resolution,
        hidden_features=hidden_features,
        hidden_layers=hidden_layers,
        output_activation=output_activation,
    )

    self.resolution = resolution
    self.rank = rank
    self.mode = mode

F-INR

finr

F-INR: a functional tensor decomposition, one SIREN per input axis.

FINR

FINR(in_features: int, out_features: int, rank: int, hidden_features: int = 256, hidden_layers: int = 3, mode: str = 'cp', first_omega: float = 100.0, hidden_omega: float = 30.0, positional_encoding: int = 0, grid: bool = True, output_activation: Callable[[Tensor], Tensor] | None = None)

Functional tensor decomposition with one SIREN per coordinate axis.

Each axis's network maps its coordinate alone to factors, and a tensor format combines the axes' factors per output channel, with no decoder after it:

  • "cp": \(y_d = \sum_r \prod_c a_{c,dr}\), for any C.
  • "tt": the tensor train \(y_d = \boldsymbol{u}_d^\top \boldsymbol{G}_d \boldsymbol{v}_d\), for C = 3, whose R x R core comes from the first axis's network, as in the authors' code.

Each network has hidden_layers sine layers and a linear head.

Parameters:

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • rank (int) –

    Rank R of the decomposition.

  • hidden_features (int, default: 256 ) –

    Hidden width.

  • hidden_layers (int, default: 3 ) –

    Number of hidden layers.

  • mode (str, default: 'cp' ) –

    Tensor format, "cp", or "tt" for C = 3.

  • first_omega (float, default: 100.0 ) –

    Frequency multiplier of the first sine layer.

  • hidden_omega (float, default: 30.0 ) –

    Frequency multiplier of the later sine layers.

  • positional_encoding (int, default: 0 ) –

    Size n of the authors' pos_enc, which replaces each network's raw coordinate by sin(2^k x) and then cos(2^k x) for k in range(int(-(n - 1) / 2), int((n + 1) / 2)), so n = 10 gives the nine bands 2^-4, ..., 2^4. 0 keeps the raw coordinate.

  • grid (bool, default: True ) –

    Whether to run each network once per distinct coordinate, when at most half of a batch's are distinct, and gather the factors per point; the output is the same. Turn it off for scattered points, such as a radiance field's samples, where finding the distinct ones only costs time.

  • output_activation (Callable[[Tensor], Tensor] | None, default: None ) –

    Callable applied to the output; None uses the identity.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).
References

Vemuri et al., "F-INR: Functional Tensor Decomposition for Implicit Neural Representations", WACV 2026.

Source code in inro/models/finr.py
def __init__(
    self,
    in_features: int,
    out_features: int,
    rank: int,
    hidden_features: int = 256,
    hidden_layers: int = 3,
    mode: str = "cp",
    first_omega: float = 100.0,
    hidden_omega: float = 30.0,
    positional_encoding: int = 0,
    grid: bool = True,
    output_activation: Callable[[Tensor], Tensor] | None = None,
) -> None:
    super().__init__()
    if mode not in ("cp", "tt"):
        raise ValueError(f"mode must be 'cp' or 'tt', got {mode!r}")
    if mode == "tt" and in_features != 3:
        raise ValueError("the tensor train combines exactly three inputs")
    if hidden_layers < 1:
        raise ValueError("hidden_layers must be at least one")

    self.in_features = in_features
    self.out_features = out_features
    self.rank = rank
    self.mode = mode
    self.hidden_omega = hidden_omega
    self.grid = grid
    self.output_activation = output_activation

    bands = torch.empty(0)
    if positional_encoding > 0:
        low = int(-(positional_encoding - 1) / 2)  # int rounds toward zero
        bands = 2.0 ** torch.arange(low, int((positional_encoding + 1) / 2))
    self.register_buffer("bands", bands, persistent=False)
    encoded = 2 * len(bands) if len(bands) else 1

    sizes = [out_features * rank] * in_features
    if mode == "tt":
        sizes[0] = out_features * rank * rank
    self.axes = nn.ModuleList(
        nn.Sequential(
            SineLayer(encoded, hidden_features, first_omega),
            *(
                SineLayer(hidden_features, hidden_features, hidden_omega)
                for _ in range(hidden_layers - 1)
            ),
            nn.Linear(hidden_features, size),
        )
        for size in sizes
    )
    self.reset_parameters()

reset_parameters

reset_parameters() -> None

Initialize as the authors' image code does, biases at zero.

Weights are drawn from U(-1/n, 1/n) in the first layer and from U(-sqrt(6/n) / hidden_omega, sqrt(6/n) / hidden_omega) after it, for fan-in n.

Source code in inro/models/finr.py
@torch.no_grad()
def reset_parameters(self) -> None:
    """Initialize as the authors' image code does, biases at zero.

    Weights are drawn from ``U(-1/n, 1/n)`` in the first layer and from
    ``U(-sqrt(6/n) / hidden_omega, sqrt(6/n) / hidden_omega)`` after it, for fan-in
    ``n``.
    """
    for network in self.axes:
        for index, layer in enumerate(network):
            linear = layer.linear if isinstance(layer, SineLayer) else layer
            fan_in = linear.in_features
            bound = (
                1 / fan_in
                if index == 0
                else math.sqrt(6 / fan_in) / self.hidden_omega
            )
            linear.weight.uniform_(-bound, bound)
            linear.bias.zero_()

Grid INR

grid_inr

Grid of independent implicit neural representations (INRs) with blended overlaps.

GridINR

GridINR(inr_class: type[Module], in_features: int, out_features: int, patches_per_axis: int | Sequence[int] = 2, overlap: float = 0.25, **inr_kwargs: Any)

Independent coordinate networks on a regular grid of patches.

Coordinates in [-1, 1] are mapped to each patch's local [-1, 1] domain. Overlapping patches blend their outputs with Hermite smoothstep weights, so up to 2**C networks contribute to each point.

Parameters:

  • inr_class (type[Module]) –

    Module class, built once per patch as inr_class(in_features, out_features, **inr_kwargs).

  • in_features (int) –

    Number of coordinate axes C.

  • out_features (int) –

    Number of output channels D.

  • patches_per_axis (int | Sequence[int], default: 2 ) –

    Number of patches, shared or per axis.

  • overlap (float, default: 0.25 ) –

    Fraction of a cell width added on each side of a patch. 0 gives hard boundaries; values above 0.5 cause discontinuities.

  • **inr_kwargs (Any, default: {} ) –

    Keyword arguments of each patch network.

Attributes:

  • inrs –

    Patch networks in row-major order, the last axis varying fastest.

Shape
  • Input: (..., C) coordinates in [-1, 1].
  • Output: (..., D).

Examples:

>>> from inro import SIREN
>>> model = GridINR(SIREN, 2, 3, patches_per_axis=2, hidden_features=16)
>>> model(torch.rand(8, 2) * 2 - 1).shape
torch.Size([8, 3])
Source code in inro/models/grid_inr.py
def __init__(
    self,
    inr_class: type[nn.Module],
    in_features: int,
    out_features: int,
    patches_per_axis: int | Sequence[int] = 2,
    overlap: float = 0.25,
    **inr_kwargs: Any,
) -> None:
    super().__init__()
    self.in_features = in_features
    self.out_features = out_features
    self.overlap = overlap

    if isinstance(patches_per_axis, int):
        patches_per_axis = (patches_per_axis,) * in_features
    if len(patches_per_axis) != in_features:
        raise ValueError(
            f"patches_per_axis must have {in_features} entries, "
            f"got {len(patches_per_axis)}"
        )
    self.patches_per_axis = patches_per_axis

    num_patches = math.prod(patches_per_axis)
    self.inrs = nn.ModuleList(
        [
            inr_class(in_features, out_features, **inr_kwargs)
            for _ in range(num_patches)
        ]
    )

    # Row-major strides map a cell coordinate to its patch network.
    strides: list[int] = []
    stride = 1
    for patch_count in reversed(patches_per_axis):
        strides.append(stride)
        stride *= patch_count
    self.register_buffer(
        "_strides", torch.tensor(list(reversed(strides)), dtype=torch.long)
    )
    self.register_buffer("_grid", torch.tensor(patches_per_axis, dtype=torch.float))
    self.register_buffer(
        "_grid_long", torch.tensor(patches_per_axis, dtype=torch.long)
    )

Deep Image Prior

dip

Deep Image Prior networks.

DIPSkip is the skip-connected encoder-decoder baseline and DIPUNet a U-Net, both adapted from https://github.com/DmitryUlyanov/deep-image-prior/tree/master/models.

References

Ulyanov et al., "Deep Image Prior", CVPR 2018.

DIPUNet

DIPUNet(in_channels: int = 16, out_channels: int = 3, feature_scale: int = 4, norm_layer: Callable[[int], Module] | None = BatchNorm2d, upsample_mode: Literal['deconv', 'nearest', 'bilinear'] = 'bilinear', padding_mode: Literal['zeros', 'reflect'] = 'reflect', bias: bool = True, sigmoid_output: bool = True)

U-Net of Deep Image Prior, with four downsampling and upsampling stages.

Parameters:

  • in_channels (int, default: 16 ) –

    Input noise channels; match the dataset's noise_channels.

  • out_channels (int, default: 3 ) –

    Output channels.

  • feature_scale (int, default: 4 ) –

    Divisor of the base widths [64, 128, 256, 512, 1024].

  • norm_layer (Callable[[int], Module] | None, default: BatchNorm2d ) –

    Normalization constructor of the input and encoder blocks; None disables it. Decoder blocks are unnormalized.

  • upsample_mode (Literal['deconv', 'nearest', 'bilinear'], default: 'bilinear' ) –

    Upsampling, "deconv" for a transposed convolution, or "nearest" or "bilinear" for interpolation followed by a convolution.

  • padding_mode (Literal['zeros', 'reflect'], default: 'reflect' ) –

    Padding, "zeros" or "reflect".

  • bias (bool, default: True ) –

    Whether to include convolution biases; transposed convolutions always do.

  • sigmoid_output (bool, default: True ) –

    Whether to apply a sigmoid after the final convolution.

Shape
  • Input: (B, in_channels, H, W).
  • Output: (B, out_channels, H', W'), with sides cropped to multiples of 16.
Source code in inro/models/dip.py
def __init__(
    self,
    in_channels: int = 16,
    out_channels: int = 3,
    feature_scale: int = 4,
    norm_layer: Callable[[int], nn.Module] | None = nn.BatchNorm2d,
    upsample_mode: Literal["deconv", "nearest", "bilinear"] = "bilinear",
    padding_mode: Literal["zeros", "reflect"] = "reflect",
    bias: bool = True,
    sigmoid_output: bool = True,
) -> None:
    super().__init__()
    widths = [width // feature_scale for width in (64, 128, 256, 512, 1024)]

    self.start = _UNetConv2(in_channels, widths[0], norm_layer, bias, padding_mode)
    self.down1 = _UNetDown(widths[0], widths[1], norm_layer, bias, padding_mode)
    self.down2 = _UNetDown(widths[1], widths[2], norm_layer, bias, padding_mode)
    self.down3 = _UNetDown(widths[2], widths[3], norm_layer, bias, padding_mode)
    self.down4 = _UNetDown(widths[3], widths[4], norm_layer, bias, padding_mode)

    self.up4 = _UNetUp(widths[3], upsample_mode, bias, padding_mode)
    self.up3 = _UNetUp(widths[2], upsample_mode, bias, padding_mode)
    self.up2 = _UNetUp(widths[1], upsample_mode, bias, padding_mode)
    self.up1 = _UNetUp(widths[0], upsample_mode, bias, padding_mode)

    self.final = _conv2d(
        widths[0], out_channels, 1, bias=bias, padding_mode=padding_mode
    )
    if sigmoid_output:
        self.final = nn.Sequential(self.final, nn.Sigmoid())

DIPSkip

DIPSkip(in_channels: int = 32, out_channels: int = 3, down_channels: Sequence[int] = (128, 128, 128, 128, 128), up_channels: Sequence[int] = (128, 128, 128, 128, 128), skip_channels: Sequence[int] = (4, 4, 4, 4, 4), down_kernel_size: int | Sequence[int] = 3, up_kernel_size: int | Sequence[int] = 3, skip_kernel_size: int = 1, up_conv1x1: bool = True, upsample_mode: str | Sequence[str] = 'bilinear', padding_mode: Literal['zeros', 'reflect'] = 'reflect', bias: bool = True, sigmoid_output: bool = True)

Deep Image Prior encoder-decoder with a skip branch at each scale.

Strided convolutions downsample the input; the decoder upsamples and concatenates skip features, center-cropped at odd sizes. Hidden convolutions use BatchNorm and LeakyReLU(0.2).

Parameters:

  • in_channels (int, default: 32 ) –

    Input noise channels; match the dataset's noise_channels.

  • out_channels (int, default: 3 ) –

    Output channels.

  • down_channels (Sequence[int], default: (128, 128, 128, 128, 128) ) –

    Encoder widths, one per scale.

  • up_channels (Sequence[int], default: (128, 128, 128, 128, 128) ) –

    Decoder widths, one per scale.

  • skip_channels (Sequence[int], default: (4, 4, 4, 4, 4) ) –

    Skip widths, one per scale; the three channel sequences must have equal length.

  • down_kernel_size (int | Sequence[int], default: 3 ) –

    Encoder kernel size, shared or per scale.

  • up_kernel_size (int | Sequence[int], default: 3 ) –

    Decoder kernel size, shared or per scale.

  • skip_kernel_size (int, default: 1 ) –

    Skip convolution kernel size.

  • up_conv1x1 (bool, default: True ) –

    Whether to add a 1x1 convolution, BatchNorm and LeakyReLU block to each decoder.

  • upsample_mode (str | Sequence[str], default: 'bilinear' ) –

    Interpolation mode, shared or per scale.

  • padding_mode (Literal['zeros', 'reflect'], default: 'reflect' ) –

    Padding, "zeros" or "reflect".

  • bias (bool, default: True ) –

    Whether to include convolution biases.

  • sigmoid_output (bool, default: True ) –

    Whether to apply a sigmoid after the final convolution.

Shape
  • Input: (B, in_channels, H, W).
  • Output: (B, out_channels, H, W); even kernel sizes shrink it.
Source code in inro/models/dip.py
def __init__(
    self,
    in_channels: int = 32,
    out_channels: int = 3,
    down_channels: Sequence[int] = (128, 128, 128, 128, 128),
    up_channels: Sequence[int] = (128, 128, 128, 128, 128),
    skip_channels: Sequence[int] = (4, 4, 4, 4, 4),
    down_kernel_size: int | Sequence[int] = 3,
    up_kernel_size: int | Sequence[int] = 3,
    skip_kernel_size: int = 1,
    up_conv1x1: bool = True,
    upsample_mode: str | Sequence[str] = "bilinear",
    padding_mode: Literal["zeros", "reflect"] = "reflect",
    bias: bool = True,
    sigmoid_output: bool = True,
) -> None:
    super().__init__()
    num_scales = len(down_channels)
    if not num_scales == len(up_channels) == len(skip_channels):
        raise ValueError(
            "down_channels, up_channels and skip_channels must have equal length"
        )
    self.num_scales = num_scales

    if not isinstance(down_kernel_size, (list, tuple)):
        down_kernel_size = [down_kernel_size] * num_scales
    if not isinstance(up_kernel_size, (list, tuple)):
        up_kernel_size = [up_kernel_size] * num_scales
    if not isinstance(upsample_mode, (list, tuple)):
        upsample_mode = [upsample_mode] * num_scales

    # Channels entering each scale from the input side and the deeper side.
    encoder_in = [in_channels, *down_channels[:-1]]
    concat_channels = [
        skip + deeper
        for skip, deeper in zip(
            skip_channels, [*up_channels[1:], down_channels[-1]]
        )
    ]
    # Modules are built group by group (all encoders, then all skips, ...)
    # so seeded initializations match the reference implementation.
    self.encoders = nn.ModuleList(
        nn.Sequential(
            *_conv2d_bn_lrelu(c_in, c_out, k, 2, bias, padding_mode),
            *_conv2d_bn_lrelu(c_out, c_out, k, 1, bias, padding_mode),
        )
        for c_in, c_out, k in zip(encoder_in, down_channels, down_kernel_size)
    )
    self.skip_convs = nn.ModuleList(
        nn.Sequential(
            *_conv2d_bn_lrelu(c_in, c_out, skip_kernel_size, 1, bias, padding_mode)
        )
        for c_in, c_out in zip(encoder_in, skip_channels)
    )
    self.upsamplers = nn.ModuleList(
        nn.Upsample(scale_factor=2, mode=mode) for mode in upsample_mode
    )
    self.concat_bns = nn.ModuleList(nn.BatchNorm2d(c) for c in concat_channels)
    self.decoders = nn.ModuleList(
        self._decoder(c_in, c_out, k, up_conv1x1, bias, padding_mode)
        for c_in, c_out, k in zip(concat_channels, up_channels, up_kernel_size)
    )

    self.final = _conv2d(
        up_channels[0], out_channels, 1, bias=bias, padding_mode=padding_mode
    )
    if sigmoid_output:
        self.final = nn.Sequential(self.final, nn.Sigmoid())