Models

The three public model classes are native torch.nn.Module objects and are available directly from nif. Each call accepts a floating-point matrix with columns ordered as [parameters, spatial coordinates] and returns a matrix of field values. Use an explicit PyTorch training loop and add model.regularization_loss(features) to its objective.

import torch
from nif import NIF

shape = dict(input_dim=1, output_dim=1, units=20, nlayers=2, activation="tanh")
parameter = dict(input_dim=1, latent_dim=2, units=10, nlayers=2, activation="tanh")
model = NIF(shape, parameter)
features = torch.tensor([[0.1, -0.2], [0.3, 0.4]])
targets = torch.zeros(2, 1)
prediction = model(features)
loss = (prediction - targets).square().mean() + model.regularization_loss(features)
loss.backward()

Configuration and precision

cfg_shape_net sets input_dim, output_dim, units and nlayers. cfg_parameter_net sets input_dim, latent_dim, units and nlayers. Plain NIF also reads each network’s activation.

Multiscale models additionally use ShapeNet omega_0, use_resblock and connectivity. Use connectivity="full" for NIFMultiScale and connectivity="last_layer" for NIFMultiScaleLastLayerParameterized. Their ParameterNet configuration includes activation and use_resblock; activation="sine" also requires omega_0. The full configurations appear in the corresponding native tutorial notebooks.

The default mixed_policy="float32" uses float32 parameters and computation. mixed_float16 and mixed_bfloat16 use autocast computation with float32 trainable parameters. Native model.to(...), model.float() and model.double() control parameter device and dtype.

Parameter-conditioned MLP field

class nif.model.NIF(cfg_shape_net, cfg_parameter_net, mixed_policy='float32')

Parameter-conditioned implicit field with shortcut MLP shape network.

Input columns are [parameters, coordinates]. build() and model() return this module. Call regularization_loss(batch) explicitly in the training objective; there are no cached batch losses or implicit fit loops.

property variable_Dtype
property compute_Dtype
forward(inputs, training=None, mask=None)

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.

call(inputs, training=None, mask=None)
build()
model()
model_p_to_w()
model_p_to_lr()
model_lr_to_w()
model_x_to_u_given_w()
regularization_loss(inputs=None)

Differentiable scalar: parameter sums, activity sum/batch, Jacobian mean.

Input-dependent regularization always uses the supplied batch. Calling this without inputs when such a penalty is configured raises ValueError. Shape-net penalties in the last-layer model use the shape coefficients.

predict(inputs, batch_size=None)

Convenience inference returning a tensor on this module’s device.

save_config(filename='config.json')

model_p_to_w() maps parameters to generated ShapeNet weights. model_p_to_lr() extracts the latent representation, while model_lr_to_w() maps that representation to generated weights. The view returned by model_x_to_u_given_w() accepts (coordinates, weights). These views share the parent model’s layer objects and gradient connections.

save_config() stores construction arguments as JSON. Save trainable tensors with torch.save(model.state_dict(), path) and restore them into a model constructed from that configuration.

Set model.eval() when switching to evaluation mode. predict() runs under torch.no_grad() and returns a tensor on the model’s device. The retained training and mask call arguments are compatibility parameters; native module mode is controlled by train() and eval().

Multiscale SIREN field

class nif.model.NIFMultiScale(cfg_shape_net, cfg_parameter_net, mixed_policy='float32')

SIREN shape network with optional two-sine residual blocks.

This class shares the training, regularization, checkpoint and subnetwork methods documented for nif.model.NIF. Its ShapeNet uses sinusoidal layers and optional two-sine residual blocks.

Linear spatial representation

class nif.model.NIFMultiScaleLastLayerParameterized(cfg_shape_net, cfg_parameter_net, mixed_policy='float32')

Linear representation u(p,x) = phi(x) a(p) + b, with learned bases.

model_p_to_lr()
model_lr_to_w()
model_x_to_phi()

ParameterNet produces the coefficients of a shared spatial basis. The model_x_to_phi() view returns basis values indexed by [batch, output component, latent feature]. model_lr_to_w() is an identity view for this architecture. Other common methods are inherited from nif.model.NIF.