Layers and derivatives

These public helpers are exported by nif.layers. They use PyTorch tensors and retain autograd connections for derivative-based training.

Pointwise derivatives

For a pointwise model with independent batch rows, the selected Jacobian has shape [batch, selected outputs, selected inputs]. The Hessian adds a second selected-input axis. y_index selects output columns; x_index selects input columns.

import torch
from nif.layers import JacobianLayer, HessianLayer

field = torch.nn.Linear(2, 1)
features = torch.tensor([[0.1, -0.2], [0.3, 0.4]])
output, jacobian = JacobianLayer(field, y_index=[0], x_index=[1])(features)
output, jacobian, hessian = HessianLayer(field, y_index=[0], x_index=[1])(features)
assert jacobian.shape == (2, 1, 1)
assert hessian.shape == (2, 1, 1, 1)
class nif.layers.JacobianLayer(model, y_index, x_index, **kwargs)

Wrap a native pointwise model and expose its selected input derivatives.

forward(x, **kwargs)

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.

get_config()
class nif.layers.HessianLayer(model, y_index, x_index, **kwargs)

Expose full output plus selected first and second input derivatives.

forward(x, **kwargs)

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.

The following function forms expose the same calculations. Note that their positional index arguments are ordered x_index, y_index.

nif.layers.compute_output_and_grad(model, x, x_index, y_index)

Return full model output and selected [batch, output, input] Jacobian.

nif.layers.compute_output_and_grad_and_hessian(model, x, x_index, y_index)

Return full output, selected Jacobian, and selected-input Hessian.

Hessian axes are [batch, selected_outputs, selected_inputs, selected_inputs].

nif.layers.compute_output_and_augment_grad(model, x, x_index, y_index)

Differentiate the second output of an augmented (output, latent) model.

Latent and activation penalties

These wrappers expect an augmented model returning (output, latent) or (output, parameter_output). Call their regularization_loss(inputs) explicitly in the objective for the current batch.

class nif.layers.JacRegLatentLayer(model, y_index, x_index, l1=0.01, mixed_policy='float32', **kwargs)

Pass through model output and explicitly penalize its latent Jacobian.

model(x) returns (output, latent). The historical l1 name is retained, but the penalty is l1 * mean(Jacobian ** 2), not an absolute norm. Call regularization_loss(inputs) in the training objective. No batch is cached, so evaluating unrelated data cannot silently change this penalty.

forward(x, **kwargs)

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.

regularization_loss(inputs=None)
get_config()
class nif.layers.ParameterOutputL1ActReg(model, l1=0.1)

Pass through augmented output and penalize all parameter activations.

model(inputs) returns (output, parameter_output). The original reduction is the sum of absolute activations over the entire batch; this wrapper intentionally does not divide that sum by the number of samples.

forward(x, **kwargs)

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.

regularization_loss(inputs=None)
get_config()

SIREN and hypernetwork layers

Custom kernels use [input, output] storage. In sinusoidal layers the argument is omega_0 * (x @ w) + b: the bias is not multiplied by omega_0.

class nif.layers.SIREN(num_inputs, num_outputs, layer_position, omega_0, cfg_shape_net=None, kernel_regularizer=None, bias_regularizer=None, mixed_policy='float32', **kwargs)

A sinusoidal layer, or a linear bottleneck/hypernetwork output layer.

mixed_policy accepts a torch dtype, a dtype name, mixed_float16, mixed_bfloat16, or an object exposing compute_dtype/variable_dtype. Explicit regularization_loss returns kernel and bias penalties.

forward(x, **kwargs)

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.

regularization_loss(inputs=None)
get_config()
get_prunable_weights()
class nif.layers.SIREN_ResNet(num_inputs, num_outputs, omega_0, kernel_regularizer=None, bias_regularizer=None, mixed_policy='float32', **kwargs)

Two sine layers with the original 0.5 * (input + hidden) residual.

Both sets start with identical samples, as in the scientific baseline, but are distinct trainable parameters.

forward(x, 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.

get_prunable_weights()
get_config()
class nif.layers.HyperLinearForSIREN(num_inputs, num_outputs, cfg_shape_net, mixed_policy='float32', connectivity='full', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, **kwargs)

Linear hypernetwork head with shape-aware SIREN initialization.

forward(x, **kwargs)

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.

regularization_loss(inputs=None)
get_prunable_weights()
get_config()
nif.layers.compute_number_of_weightbias_by_its_position_for_shapenet(cfg_shape_net)

Count generated first, hidden and final kernels (excluding biases).

nif.layers.gen_hypernetwork_weights_bias_for_siren_shapenet(num_inputs, num_outputs, weight_factor, num_weight_first, num_weight_hidden, num_weight_last, input_dim, width, omega_0, variable_dtype)

Generate hypernetwork kernels and shape-aware initial output biases.

Generated entries have first-layer, hidden-layer, final-layer kernels, followed by all biases. Last-layer kernels use Glorot scaling, since they have no sine activation; all generated biases use 1 / width scaling.

Dense and residual helpers

class nif.layers.Dense(in_features, out_features, bias=True, device=None, dtype=None)

A direct alias of torch.nn.Linear. Use the native PyTorch constructor and its weight layout [output, input]. See the PyTorch Linear reference.

class nif.layers.MLP_ResNet(width, activation, kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, mixed_policy='float32', **kwargs)

Two-layer residual block: activation(x + L2(activation(L1(x)))).

forward(x, **kwargs)

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.

regularization_loss(inputs=None)
get_prunable_weights()
get_config()
class nif.layers.MLP_SimpleShortCut(width, activation, kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, mixed_policy='float32', **kwargs)

One activated dense layer with a direct input shortcut: x + L1(x).

forward(x, **kwargs)

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.

regularization_loss(inputs=None)
get_prunable_weights()
get_config()
class nif.layers.EinsumLayer(equation, **kwargs)

Wrap torch.einsum; call with a tuple of operand tensors.

forward(inputs, *args, **kwargs)

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.

get_config()
class nif.layers.BiasAddLayer(output_dim, mixed_policy='float32', **kwargs)

Add a trainable bias initialized from N(0,.1), truncated at ±.2.

forward(inputs)

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.

get_config()