Networks

networks.energies module

Models for energy functions.

class networks.energies.MLP(dim_hidden: ~typing.Sequence[int], act_fn: ~typing.Callable = <PjitFunction of <function softplus>>, parent: ~typing.Type[~flax.linen.module.Module] | ~flax.core.scope.Scope | ~typing.Type[~flax.linen.module._Sentinel] | None = <flax.linen.module._Sentinel object>, name: str | None = None)[source]

Bases: Module

Simple energy model.

Source: https://github.com/bunnech/jkonet

act_fn() → Array

Softplus activation function.

Computes the element-wise function

\[\mathrm{softplus}(x) = \log(1 + e^x)\]
Parameters:

x – input array

dim_hidden: Sequence[int]
name: str | None = None
parent: Type[Module] | Scope | Type[_Sentinel] | None = None
scope: Scope | None = None
setup()[source]

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

  1. Immediately when invoking apply(), init() or init_and_output().

  2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

    >>> class MyModule(nn.Module):
...   def setup(self):
...     submodule = nn.Conv(...)

...     # Accessing `submodule` attributes does not yet work here.

...     # The following line invokes `self.__setattr__`, which gives
...     # `submodule` the name "conv1".
...     self.conv1 = submodule

...     # Accessing `submodule` attributes or methods is now safe and
...     # either causes setup() to be called once.

  3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

networks.fixpoint_loop module

This module provides a backprop-friendly fixed point loop implementation for JAX.

Authors: Jonathan Heek, Marco Cuturi, Charlotte Bunne Source: https://github.com/bunnech/jkonet

networks.fixpoint_loop.fixpoint_iter_bwd(cond_fn: Callable[[int, Any, Any], bool], body_fn: Callable[[int, Any, Any, bool], Any], min_iterations: int, max_iterations: int, inner_iterations: int, res: Tuple[Any, int, Any], g: Any) → Tuple[Any, Any][source]

Backward pass for the fixed-point iteration loop.

Parameters:

  • cond_fn (Callable[[int, Any, Any], bool]) – A function that was used in the forward pass to determine whether the loop should continue.

  • body_fn (Callable[[int, Any, Any, bool], Any]) – A function that defines the operations performed in each iteration during the forward pass.

  • min_iterations (int) – The minimum number of iterations that was performed in the forward pass.

  • max_iterations (int) – The maximum number of iterations that was performed in the forward pass.

  • inner_iterations (int) – The number of iterations in each inner loop during the forward pass.

  • res (Tuple[Any, int, Any]) – A tuple containing the constants, final iteration count, and recorded intermediate states from the forward pass.

  • g (Any) – The gradient with respect to the final state.

Returns:

A tuple containing the gradients with respect to the constants and the initial state.

Return type:

Tuple[Any, Any]

networks.fixpoint_loop.fixpoint_iter_fwd(cond_fn: Callable[[int, Any, Any], bool], body_fn: Callable[[int, Any, Any, bool], Any], min_iterations: int, max_iterations: int, inner_iterations: int, constants: Any, state: Any) → Tuple[Any, Tuple[Any, int, Any]][source]

Forward pass for the fixed-point iteration loop, storing intermediate states.

Parameters:

  • cond_fn (Callable[[int, Any, Any], bool]) – A function that determines whether the loop should continue based on the current iteration, constants, and state.

  • body_fn (Callable[[int, Any, Any, bool], Any]) – A function that defines the operations to perform in each iteration. It takes the current iteration number, constants, state, and a boolean indicating whether to compute an error.

  • min_iterations (int) – Lower bound on the total number of fixed point iterations.

  • max_iterations (int) – Upper bound on the total number of fixed point iterations.

  • inner_iterations (int) – Default number of iterations in the inner loop.

  • constants (Any) – Constant parameters passed to the body function during the loop.

  • state (Any) – The initial state of the loop.

Returns:

The final state after the loop terminates, and a tuple containing the constants, the final iteration number, and the recorded intermediate states.

Return type:

Tuple[Any, Tuple[Any, int, Any]]

networks.icnns module

This module contains the implementation of the Input Convex Neural Network (ICNN) model.

Source: https://github.com/bunnech/jkonet

class networks.icnns.Dense(dim_hidden: int, beta: float = 1.0, use_bias: bool = True, dtype: Any = <class 'jax.numpy.float32'>, precision: Any = None, kernel_init: Callable[[Any, Tuple[int], Any], Any] = <function variance_scaling.<locals>.init at 0x79afb9706b60>, bias_init: Callable[[Any, Tuple[int], Any], Any] = <function zeros at 0x79b076036b60>, parent: Union[Type[flax.linen.module.Module], flax.core.scope.Scope, Type[flax.linen.module._Sentinel], NoneType] = <flax.linen.module._Sentinel object at 0x79afbb953cb0>, name: Optional[str] = None)[source]

Bases: Module

beta: float = 1.0
bias_init(shape: ~collections.abc.Sequence[int | ~typing.Any], dtype: ~typing.Any = <class 'jax.numpy.float64'>) → Array

An initializer that returns a constant array full of zeros.

The key argument is ignored.

>>> import jax, jax.numpy as jnp
>>> jax.nn.initializers.zeros(jax.random.key(42), (2, 3), jnp.float32)
Array([[0., 0., 0.],
       [0., 0., 0.]], dtype=float32)
dim_hidden: int
dtype

alias of float32

kernel_init(shape: ~collections.abc.Sequence[int | ~typing.Any], dtype: ~typing.Any = <class 'jax.numpy.float64'>) → Array
name: str | None = None
parent: Type[Module] | Scope | Type[_Sentinel] | None = None
precision: Any = None
scope: Scope | None = None
use_bias: bool = True
class networks.icnns.ICNN(dim_hidden: Sequence[int], init_std: float = 0.1, init_fn: str = 'normal', act_fn: Callable = <PjitFunction of <function leaky_relu at 0x79b0760def20>>, pos_weights: bool = True, parent: Union[Type[flax.linen.module.Module], flax.core.scope.Scope, Type[flax.linen.module._Sentinel], NoneType] = <flax.linen.module._Sentinel object at 0x79afbb953cb0>, name: Optional[str] = None)[source]

Bases: Module

act_fn(negative_slope: Array | ndarray | bool_ | number | bool | int | float | complex = 0.01) → Array

Leaky rectified linear unit activation function.

Computes the element-wise function:

\[\begin{split}\mathrm{leaky\_relu}(x) = \begin{cases} x, & x \ge 0\\ \alpha x, & x < 0 \end{cases}\end{split}\]

where \(\alpha\) = negative_slope.

Parameters:

  • x – input array

  • negative_slope – array or scalar specifying the negative slope (default: 0.01)

Returns:

An array.

See also

relu()

dim_hidden: Sequence[int]
init_fn: str = 'normal'
init_std: float = 0.1
name: str | None = None
parent: Type[Module] | Scope | Type[_Sentinel] | None = None
pos_weights: bool = True
scope: Scope | None = None
setup()[source]

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

  1. Immediately when invoking apply(), init() or init_and_output().

  2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

    >>> class MyModule(nn.Module):
...   def setup(self):
...     submodule = nn.Conv(...)

...     # Accessing `submodule` attributes does not yet work here.

...     # The following line invokes `self.__setattr__`, which gives
...     # `submodule` the name "conv1".
...     self.conv1 = submodule

...     # Accessing `submodule` attributes or methods is now safe and
...     # either causes setup() to be called once.

  3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

networks.optim module

This module contains optimization utils used to train the models.

Source: https://github.com/bunnech/jkonet

Functions

  • get_optimizer

    Returns an Optax optimizer object based on the provided configuration.

  • create_train_state

    Creates an initial TrainState for the given model and optimizer.

  • create_train_state_from_params

    Creates a TrainState from existing model parameters.

  • global_norm

    Computes the global norm of gradients across a nested structure of tensors.

  • clip_weights_icnn

    Clip the weights of an Input Convex Neural Network (ICNN).

networks.optim.clip_weights_icnn(params: FrozenDict) → FrozenDict[source]

Clip the weights of an Input Convex Neural Network (ICNN).

This function modifies the weights of the ICNN by clipping the values in kernels that start with β€˜Wz’ to ensure they are non-negative. This is necessary to maintain the convexity property of the ICNN.

Parameters:

params (FrozenDict) – A frozen dictionary containing the parameters of the ICNN.

Returns:

A frozen dictionary with the same structure as params, but with the relevant weights clipped to be non-negative.

Return type:

Any

networks.optim.create_train_state(rng: PRNGKey, model: Module, optimizer: GradientTransformation, input_shape: int) → TrainState[source]

Creates an initial TrainState for the given model and optimizer.

Parameters:

  • rng (jax.random.PRNGKey) – Random key used for initializing the model parameters.

  • model (nn.Module) – Flax model used for creating the initial state.

  • optimizer (optax.GradientTransformation) – Optimizer object used for updating the model parameters.

  • input_shape (int) – Shape of the input data used to initialize the model.

Returns:

The initialized train state containing model parameters and optimizer.

Return type:

train_state.TrainState

networks.optim.create_train_state_from_params(model: Module, params: Dict[str, Any], optimizer: GradientTransformation) → TrainState[source]

Creates a TrainState from existing model parameters.

Parameters:

  • model (nn.Module) – Flax model used for creating the initial state.

  • params (Dict[str, Any]) – Dictionary of model parameters.

  • optimizer (optax.GradientTransformation) – Optimizer object used for updating the model parameters.

Returns:

The train state containing the provided model parameters and optimizer.

Return type:

train_state.TrainState

networks.optim.get_optimizer(config: Dict[str, Any]) → GradientTransformation[source]

Returns an Optax optimizer object based on the provided configuration.

Parameters:

config (Dict[str, Any]) –

Dictionary containing optimizer configuration. Expected keys are:

  • ’optimizer’: The name of the optimizer (β€˜Adam’ or β€˜SGD’).

  • ’lr’: Learning rate for the optimizer.

  • ’beta1’: Beta1 parameter for the Adam optimizer.

  • ’beta2’: Beta2 parameter for the Adam optimizer.

  • ’eps’: Epsilon parameter for the Adam optimizer.

  • ’grad_clip’: Optional maximum global norm for gradient clipping.

Returns:

The configured Optax optimizer object.

Return type:

optax.GradientTransformation

Raises:

NotImplementedError – If the optimizer name is not supported.

networks.optim.global_norm(updates: Dict[str, Array]) → Array[source]

Computes the global norm of gradients across a nested structure of tensors.

Parameters:

updates (Dict[str, jnp.ndarray]) – Dictionary where values are tensors (e.g., gradients).

Returns:

The global norm of the gradients.

Return type:

jnp.ndarray

networks.optim.penalize_weights_icnn(params: FrozenDict) → Array[source]

Compute a penalty for negative weights in an ICNN.

This function calculates a penalty term based on the L2 norm of any negative values in the weights that start with β€˜Wz’. This penalty can be added to the loss function during training to encourage the network to maintain non-negative weights in those layers, which is important for the ICNN’s convexity.

Parameters:

params (FrozenDict) – A frozen dictionary containing the parameters of the ICNN.

Returns:

A scalar penalty value representing the sum of the L2 norms of the negative weights.

Return type:

jnp.ndarray

networks.utils module

Module for gradient computation and parameter analysis in Flax neural networks using JAX.

Functions

  • network_grad

    Computes the gradient of the network’s output with respect to its input. The gradient is evaluated for each sample using vectorized mapping (vmap).

  • network_grad_time

    Computes the gradient of the network’s output with respect to its input, excluding the time component.

  • count_parameters

    Returns the total number of parameters in the given Flax neural network model.

networks.utils.count_parameters(model: Module) → int[source]

Counts the total number of parameters in the model.

Parameters:

model (nn.Module) – The Flax neural network module.

Returns:

The total number of parameters in the model.

Return type:

int

networks.utils.network_grad(network: Module, params: Dict[str, Array]) → Callable[[Array], Array][source]

Computes the gradient of the network’s output with respect to its input for each sample.

Parameters:

  • network (nn.Module) – The Flax neural network module.

  • params (Dict[str, jnp.ndarray]) – Dictionary containing model parameters.

Returns:

A function that computes gradients with respect to the network’s input.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]

networks.utils.network_grad_time(network: Module, params: Dict[str, Array]) → Callable[[Array], Array][source]

Computes the gradient of the network’s output with respect to the input, excluding the time component.

In the time-varying JKOnet* model, the gradient in the loss is computed with respect to the input, excluding the time component.

Parameters:

  • network (nn.Module) – The Flax neural network module.

  • params (Dict[str, jnp.ndarray]) – Dictionary containing model parameters.

Returns:

A function that computes gradients with respect to the input, excluding the time component.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]