nif.layers

class nif.layers.SIREN(*args, **kwargs)

Bases: Layer, PrunableLayer

A class representing the SIREN layer.

Parameters:

num_inputs (int) – The number of input units.
num_outputs (int) – The number of output units.
layer_position (str) – The position of the SIREN layer in the network architecture. Possible values are ‘first’, ‘hidden’, ‘last’ and ‘bottleneck’.
omega_0 (float) – The cutoff frequency for the initial frequency. Should be set to 1.0 for most cases.
cfg_shape_net (dict) – A dictionary containing the configuration parameters for the network if the layer position is ‘last’.
kernel_regularizer (tf.keras.regularizers.Regularizer) – Regularizer function applied to the kernel weights matrix.
bias_regularizer (tf.keras.regularizers.Regularizer) – Regularizer function applied to the bias vector.
mixed_policy (tf.keras.mixed_precision.Policy) – A mixed precision policy used for the weights and biases.
**kwargs – Additional arguments.

Variables:

w_init (tf.Tensor) – The initialized weights.
b_init (tf.Tensor) – The initialized biases.
w (tf.Variable) – The learnable weights.
b (tf.Variable) – The learnable biases.
compute_Dtype (tf.DType) – The computational datatype of the layer.

call(inputs, **kwargs): Defines the computation performed at every call.

get_config(): Returns the config of the layer.

get_prunable_weights(): Returns the prunable weights of the layer.

call(x, **kwargs)

Compute the output of the layer given an input tensor x.

Parameters:: x (tf.Tensor) – Input tensor of shape (batch_size, input_dim).
Returns:: tf.Tensor – Output tensor of shape (batch_size, output_dim).

get_config()

Returns the configuration of the layer.

Returns:: dict – The configuration of the layer.

get_prunable_weights()

Returns the list of prunable weights of the layer.

Returns:: list – The list of prunable weights of the layer.

class nif.layers.SIREN_ResNet(*args, **kwargs)

Bases: SIREN

A subclass of the SIREN class implementing a residual block.

Parameters:

num_inputs (int) – Number of input features.
num_outputs (int) – Number of output features.
omega_0 (float) – Frequency parameter for the SIREN activation function.
kernel_regularizer (tf.keras.regularizers.Regularizer) – Regularizer function applied to the layer’s weights.
bias_regularizer (tf.keras.regularizers.Regularizer) – Regularizer function applied to the layer’s biases.
mixed_policy (tf.keras.mixed_precision.Policy) – Policy to use for mixed precision computation. Defaults to “float32”.
**kwargs – Additional keyword arguments to pass to the parent class constructor.

Variables:

w2 (tf.Variable) – Weight variable for the second layer of the residual block.
b2 (tf.Variable) – Bias variable for the second layer of the residual block.

call(x, training=None, mask=None): Performs a forward pass through the layer.

get_prunable_weights(): Returns a list of prunable weight variables.

call(x, training=None, mask=None)

Performs a forward pass through the layer.

Parameters:

x (tf.Tensor) – Input tensor.
training (bool) – Whether the layer is in training mode.
mask – Ignored.

Returns:

The output tensor of the layer.

get_prunable_weights()

Returns the list of prunable weights in the layer, i.e., the weights that can be: pruned during training.

Returns:: List of tf.Variable objects representing the prunable weights in the layer.

class nif.layers.Dense(*args, **kwargs)

Bases: Layer

Just your regular densely-connected NN layer.

Dense implements the operation: output = activation(dot(input, kernel) + bias) where activation is the element-wise activation function passed as the activation argument, kernel is a weights matrix created by the layer, and bias is a bias vector created by the layer (only applicable if use_bias is True). These are all attributes of Dense.

Note: If the input to the layer has a rank greater than 2, then Dense computes the dot product between the inputs and the kernel along the last axis of the inputs and axis 0 of the kernel (using tf.tensordot). For example, if input has dimensions (batch_size, d0, d1), then we create a kernel with shape (d1, units), and the kernel operates along axis 2 of the input, on every sub-tensor of shape (1, 1, d1) (there are batch_size * d0 such sub-tensors). The output in this case will have shape (batch_size, d0, units).

Besides, layer attributes cannot be modified after the layer has been called once (except the trainable attribute). When a popular kwarg input_shape is passed, then keras will create an input layer to insert before the current layer. This can be treated equivalent to explicitly defining an InputLayer.

Example:

>>> # Create a `Sequential` model and add a Dense layer as the first layer.
>>> model = tf.keras.models.Sequential()
>>> model.add(tf.keras.Input(shape=(16,)))
>>> model.add(tf.keras.layers.Dense(32, activation='relu'))
>>> # Now the model will take as input arrays of shape (None, 16)
>>> # and output arrays of shape (None, 32).
>>> # Note that after the first layer, you don't need to specify
>>> # the size of the input anymore:
>>> model.add(tf.keras.layers.Dense(32))
>>> model.output_shape
(None, 32)

Parameters:

units – Positive integer, dimensionality of the output space.
activation – Activation function to use. If you don’t specify anything, no activation is applied (ie. “linear” activation: a(x) = x).
use_bias – Boolean, whether the layer uses a bias vector.
kernel_initializer – Initializer for the kernel weights matrix.
bias_initializer – Initializer for the bias vector.
kernel_regularizer – Regularizer function applied to the kernel weights matrix.
bias_regularizer – Regularizer function applied to the bias vector.
activity_regularizer – Regularizer function applied to the output of the layer (its “activation”).
kernel_constraint – Constraint function applied to the kernel weights matrix.
bias_constraint – Constraint function applied to the bias vector.

Input shape:: N-D tensor with shape: (batch_size, …, input_dim). The most common situation would be a 2D input with shape (batch_size, input_dim).
Output shape:: N-D tensor with shape: (batch_size, …, units). For instance, for a 2D input with shape (batch_size, input_dim), the output would have shape (batch_size, units).

build(input_shape)

Creates the variables of the layer (optional, for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call. It is invoked automatically before the first execution of call().

This is typically used to create the weights of Layer subclasses (at the discretion of the subclass implementer).

Parameters:: input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

call(inputs)

This is where the layer’s logic lives.

The call() method may not create state (except in its first invocation, wrapping the creation of variables or other resources in tf.init_scope()). It is recommended to create state in __init__(), or the build() method that is called automatically before call() executes the first time.

Parameters:

inputs – Input tensor, or dict/list/tuple of input tensors. The first positional inputs argument is subject to special rules: - inputs must be explicitly passed. A layer cannot have zero

arguments, and inputs cannot be provided via the default value of a keyword argument.
- NumPy array or Python scalar values in inputs get cast as tensors.
- Keras mask metadata is only collected from inputs.
- Layers are built (build(input_shape) method) using shape info from inputs only.
- input_spec compatibility is only checked against inputs.
- Mixed precision input casting is only applied to inputs. If a layer has tensor arguments in *args or **kwargs, their casting behavior in mixed precision should be handled manually.
- The SavedModel input specification is generated using inputs only.
- Integration with various ecosystem packages like TFMOT, TFLite, TF.js, etc is only supported for inputs and not for tensors in positional and keyword arguments.
*args – Additional positional arguments. May contain tensors, although this is not recommended, for the reasons above.
**kwargs – Additional keyword arguments. May contain tensors, although this is not recommended, for the reasons above. The following optional keyword arguments are reserved: - training: Boolean scalar tensor of Python boolean indicating

whether the call is meant for training or inference.
- mask: Boolean input mask. If the layer’s call() method takes a mask argument, its default value will be set to the mask generated for inputs by the previous layer (if input did come from a layer that generated a corresponding mask, i.e. if it came from a Keras layer with masking support).

Returns:

A tensor or list/tuple of tensors.

compute_output_shape(input_shape)

Computes the output shape of the layer.

This method will cause the layer’s state to be built, if that has not happened before. This requires that the layer will later be used with inputs that match the input shape provided here.

Parameters:: input_shape – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns:: An input shape tuple.

get_config()

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Returns:: Python dictionary.

class nif.layers.HyperLinearForSIREN(*args, **kwargs)

Bases: Layer, PrunableLayer

Implements a hypernetwork that generates weights and biases for the SIREN layer.

Parameters:

num_inputs (int) – Number of input units.
num_outputs (int) – Number of output units.
cfg_shape_net (dict) – Configuration dictionary of the shape network.
mixed_policy (tf.keras.mixed_precision.Policy) – Policy for mixed precision computation.
connectivity (str) – Connectivity type of the SIREN layer. Should be set to full or last_layer.
kernel_regularizer (tf.keras.regularizers.Regularizer) – Regularizer for the kernel.
bias_regularizer (tf.keras.regularizers.Regularizer) – Regularizer for the bias.
activity_regularizer (tf.keras.regularizers.Regularizer) – Regularizer for the layer activity.
**kwargs – Additional layer arguments.

Variables:

kernel_regularizer (tf.keras.regularizers.Regularizer) – Regularizer for the kernel.
bias_regularizer (tf.keras.regularizers.Regularizer) – Regularizer for the bias.
compute_Dtype (tf.dtypes.DType) – Data type for computation.
w (tf.Variable) – Variable for the weights.
b (tf.Variable) – Variable for the biases.

call(self, x, **kwargs): Computes the output of the layer for the input x.

get_config(self): Returns the configuration dictionary of the layer.

get_prunable_weights(self): Returns the prunable weights of the layer.

call(x, **kwargs)

This is where the layer’s logic lives.

The call() method may not create state (except in its first invocation, wrapping the creation of variables or other resources in tf.init_scope()). It is recommended to create state in __init__(), or the build() method that is called automatically before call() executes the first time.

Parameters:

inputs – Input tensor, or dict/list/tuple of input tensors. The first positional inputs argument is subject to special rules: - inputs must be explicitly passed. A layer cannot have zero

arguments, and inputs cannot be provided via the default value of a keyword argument.
- NumPy array or Python scalar values in inputs get cast as tensors.
- Keras mask metadata is only collected from inputs.
- Layers are built (build(input_shape) method) using shape info from inputs only.
- input_spec compatibility is only checked against inputs.
- Mixed precision input casting is only applied to inputs. If a layer has tensor arguments in *args or **kwargs, their casting behavior in mixed precision should be handled manually.
- The SavedModel input specification is generated using inputs only.
- Integration with various ecosystem packages like TFMOT, TFLite, TF.js, etc is only supported for inputs and not for tensors in positional and keyword arguments.
*args – Additional positional arguments. May contain tensors, although this is not recommended, for the reasons above.
**kwargs – Additional keyword arguments. May contain tensors, although this is not recommended, for the reasons above. The following optional keyword arguments are reserved: - training: Boolean scalar tensor of Python boolean indicating

whether the call is meant for training or inference.
- mask: Boolean input mask. If the layer’s call() method takes a mask argument, its default value will be set to the mask generated for inputs by the previous layer (if input did come from a layer that generated a corresponding mask, i.e. if it came from a Keras layer with masking support).

Returns:

A tensor or list/tuple of tensors.

get_config()

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Returns:: Python dictionary.

get_prunable_weights()

Returns list of prunable weight tensors.

All the weight tensors which the layer wants to be pruned during training must be returned by this method.

Returns: List of weight tensors/kernels in the keras layer which must be: pruned during training.

class nif.layers.MLP_ResNet(*args, **kwargs)

Bases: Layer, PrunableLayer

A fully connected neural network with residual connections.

Variables:

compute_Dtype (tf.dtypes.DType) – The floating-point precision used for computation.
variable_Dtype (tf.dtypes.DType) – The floating-point precision used for variables.
act (function) – The activation function to use.
L1 (tf.keras.layers.Dense) – The first fully connected layer.
L2 (tf.keras.layers.Dense) – The second fully connected layer.

Parameters:

width (int) – The width of the fully connected layers.
activation (str) – The name of the activation function to use.
kernel_initializer (str) – The name of the initializer to use for the kernel weights.
bias_initializer (str) – The name of the initializer to use for the bias weights.
kernel_regularizer (str) – The name of the regularizer to use for the kernel weights.
bias_regularizer (str) – The name of the regularizer to use for the bias weights.
mixed_policy (tf.keras.mixed_precision.Policy) – The policy to use for mixed-precision training.
**kwargs – Additional keyword arguments to pass to the base class constructor.

call(x, **kwargs)

Forward pass of the neural network with residual connections.

Parameters:: x (tf.Tensor) – The input tensor to the network.
Returns:: tf.Tensor – The output tensor of the network.

get_config()

Returns the configuration of the layer.

Returns:: dict – The configuration of the layer.

get_prunable_weights()

Returns the list of prunable weights of the layer.

Returns:: list – The list of prunable weights of the layer.

class nif.layers.MLP_SimpleShortCut(*args, **kwargs)

Bases: Layer, PrunableLayer

A fully connected layer with a skip connection that adds the input tensor to the output tensor. Inherits from tf.keras.layers.Layer and tfmot.sparsity.keras.PrunableLayer.

Parameters:

width (int) – The number of neurons in the layer.
activation (str) – The activation function to use for the layer.
kernel_initializer – Initializer for the kernel weights matrix.
bias_initializer – Initializer for the bias vector.
kernel_regularizer – Regularizer function applied to the kernel weights matrix.
bias_regularizer – Regularizer function applied to the bias vector.
mixed_policy – Floating-point precision to use for computing the layer.
**kwargs – Additional keyword arguments to pass to the base class constructor.

call(x, **kwargs)

Applies the fully connected layer with a skip connection to the input tensor.

Parameters:

x – Input tensor to the layer.
**kwargs – Additional keyword arguments to pass to the call method.

Returns:

The result of applying the layer to the input tensor with a skip connection added.

get_config()

Returns the configuration of the layer.

Returns:: A dictionary containing the configuration of the layer.

get_prunable_weights()

Returns the weights of the layer that can be pruned.

Returns:: The weights of the layer that can be pruned.

class nif.layers.JacRegLatentLayer(*args, **kwargs)

Bases: JacobianLayer

A custom Keras layer that applies Jacobian regularization to a TensorFlow model’s output.

This layer computes the Jacobian matrix of the output with respect to the input, and adds a regularization term to the loss function to encourage the Jacobian to be small. The regularization term is given by:

L = l1 * mean((df/dx)^2)

where l1 is a hyperparameter controlling the strength of the regularization, f is the TensorFlow model, y is its output, and x is its input.

Parameters:

model (tf.keras.Model) – The TensorFlow model for which to compute the Jacobian.
y_index (int or List[int]) – The index or indices of the output variable(s) to compute the Jacobian with respect to.
x_index (int or List[int]) – The index or indices of the input variable(s) to compute the Jacobian with respect to.
l1 (float) – The weight of the Jacobian regularization term in the loss function.
mixed_policy (str) – The floating-point precision to use for computing the Jacobian.
**kwargs – Additional keyword arguments to pass to the base class constructor.

call(x, **kwargs)

Computes the output of the model and the Jacobian regularization loss.

Parameters:

x (tf.Tensor) – The input tensor(s) to the model.
**kwargs – Additional keyword arguments to pass to the underlying TensorFlow function.

Returns:

tf.Tensor – The output of the model.

get_config()

Returns the configuration of the layer.

Returns:: Dict[str, Any] – The configuration of the layer.

class nif.layers.JacobianLayer(*args, **kwargs)

Bases: Layer

A custom Keras layer that computes the Jacobian matrix of a TensorFlow model.

Parameters:

model (tf.keras.Model) – The TensorFlow model for which to compute the Jacobian.
y_index (int or List[int]) – The index or indices of the output variable(s) to compute the Jacobian with respect to.
x_index (int or List[int]) – The index or indices of the input variable(s) to compute the Jacobian with respect to.

Returns:

Tuple[tf.Tensor, tf.Tensor] – A tuple containing the output of the model and the Jacobian matrix.

call(x, **kwargs)

Computes the output of the model and the Jacobian matrix.

Parameters:

x (tf.Tensor) – The input tensor(s) to the model.
**kwargs – Additional keyword arguments to pass to the underlying TensorFlow function.

Returns:

Tuple[tf.Tensor, tf.Tensor] – A tuple containing the output of the model and the Jacobian matrix.

class nif.layers.HessianLayer(*args, **kwargs)

Bases: Layer

A custom Keras layer that computes the Hessian matrix of a TensorFlow model’s output.

This layer computes the Hessian matrix of the output with respect to the input, and returns it along with the first and second derivatives of the output with respect to the input. The first derivative is given by the Jacobian matrix, and the second derivative is the Hessian matrix.

Parameters:

model (tf.keras.Model) – The TensorFlow model for which to compute the Hessian.
y_index (int or List[int]) – The index or indices of the output variable(s) to compute the Hessian with respect to.
x_index (int or List[int]) – The index or indices of the input variable(s) to compute the Hessian with respect to.

call(x, **kwargs)

Computes the output of the model, the Jacobian matrix, and the Hessian matrix.

Parameters:

x (tf.Tensor) – The input tensor(s) to the model.
**kwargs – Additional keyword arguments to pass to the underlying TensorFlow function.

Returns:

Tuple[tf.Tensor, tf.Tensor, tf.Tensor] –

A tuple containing the output of the model,: the Jacobian matrix, and the Hessian matrix.

class nif.layers.ParameterOutputL1ActReg(*args, **kwargs)

Bases: Layer

Custom Keras layer for parameter output L1 activation regularization.

Parameters:

model (tf.keras.Model) – The TensorFlow model whose parameter output activations are to be regularized.
l1 (float) – The weight of the L1 regularization term in the loss function.

call(x, **kwargs)

Computes the output of the layer.

Parameters:: x (tf.Tensor) – The input tensor.
Returns:: The output tensor with parameter output L1 activation regularization applied.

class nif.layers.EinsumLayer(*args, **kwargs)

Bases: Layer

A custom layer that wraps a single tf.einsum operation.

Usage: x = EinsumLayer(“bmhwf,bmoh->bmowf”)((x1, x2))

Parameters:

equation (str) – The Einstein summation notation equation to use for the operation.
**kwargs – Additional keyword arguments to pass to the base class constructor.

call(inputs, *args, **kwargs)

Performs the tf.einsum operation on the inputs.

Parameters:: inputs (tuple[tf.Tensor]) – A tuple of input tensors to the operation.
Returns:: tf.Tensor – The output tensor of the operation.

get_config()

Returns the configuration of the layer.

Returns:: dict – A dictionary containing the configuration of the layer.

class nif.layers.BiasAddLayer(*args, **kwargs)

Bases: Layer

A custom layer that adds a bias vector to the inputs.

Parameters:

output_dim (int) – The dimensionality of the output space.
mixed_policy (str) – The floating-point precision to use for the bias vector.
**kwargs – Additional keyword arguments to pass to the base class constructor.

call(inputs)

Adds the bias vector to the input tensor.

Parameters:: inputs (tf.Tensor) – The input tensor to add the bias vector to.
Returns:: tf.Tensor – The input tensor with the bias vector added to it.

get_config()

Returns the configuration of the layer.

Returns:: dict – A dictionary containing the configuration of the layer.