+
+
+
+
+
+
+
+
+
+
+
+
+
+
+ Advanced users can access all the underlying functions of jpc.make_pc_step as
+well as additional features. A custom PC training step looks like the following:
+
import jpc
+
+# 1. initialise activities with a feedforward pass
+activities = jpc.init_activities_with_ffwd(model=model, input=x)
+
+# 2. run inference to equilibrium
+equilibrated_activities = jpc.solve_inference(
+ params=(model, None),
+ activities=activities,
+ output=y,
+ input=x
+)
+
+# 3. update parameters at the activities' solution with PC
+param_update_result = jpc.update_params(
+ params=(model, None),
+ activities=equilibrated_activities,
+ optim=param_optim,
+ opt_state=param_opt_state,
+ output=y,
+ input=x
+)
+
+# updated model and optimiser
+model = param_update_result["model"]
+param_optim = param_update_result["optim"]
+param_opt_state = param_update_result["opt_state"]
+activity_optim = optax.sgd(1e-3)
+
+# 1. initialise activities
+...
+
+# 2. infer with gradient descent
+activity_opt_state = activity_optim.init(activities)
+
+for t in range(T):
+ activity_update_result = jpc.update_activities(
+ params=(model, None),
+ activities=activities,
+ optim=activity_optim,
+ opt_state=activity_opt_state,
+ output=y,
+ input=x
+ )
+ # updated activities and optimiser
+ activities = activity_update_result["activities"]
+ activity_optim = activity_update_result["optim"]
+ activity_opt_state = activity_update_result["opt_state"]
+
+# 3. update parameters at the activities' solution with PC
+...
+jpc.linear_equilib_energy(network: PyTree[equinox.nn._linear.Linear], x: ArrayLike, y: ArrayLike) -> Array
+
+
+¤Computes the theoretical equilibrated PC energy for a deep linear network (DLN).
+where the rescaling is \(S = I_{d_y} + \sum_{\ell=2}^L (W_{L:\ell})(W_{L:\ell})^T\), +and we use the shorthand \(W_{k:\ell} = W_k \dots W_\ell\) for \(\ell, k \in 1,\dots, L\). +See the reference below for more details.
+Note
+This expression assumes no biases.
+@article{innocenti2024only,
+ title={Only Strict Saddles in the Energy Landscape of Predictive Coding Networks?},
+ author={Innocenti, Francesco and Achour, El Mehdi and Singh, Ryan and Buckley, Christopher L},
+ journal={arXiv preprint arXiv:2408.11979},
+ year={2024}
+}
+Main arguments:
+network: Linear network defined as a list of Equinox Linear layers.x: Network input.y: Network output.Returns:
+Mean total analytical energy over a batch or dataset.
+ +jpc.linear_activities_solution(network: PyTree[equinox.nn._linear.Linear], x: ArrayLike, y: ArrayLike) -> PyTree[Array]
+
+
+¤Computes the theoretical solution for the PC activities of a deep linear network (DLN).
+where \(A\) is a sparse block diagonal matrix depending only on the weights, +and \(\mathbf{b} = [W_1 \mathbf{x}, \mathbf{0}, \dots, W_L^T \mathbf{y}]^T\). +In particular, \(A_{\ell,k} = I + W_\ell^T W_\ell\) if \(\ell = k\), +\(A_{\ell,k} = -W_\ell\) if \(\ell = k+1\), +\(A_{\ell,k} = -W_\ell^T\) if \(\ell = k-1\), and \(\mathbf{0}\) otherwise, +for \(\ell, k \in [2, \dots, L]\).
+Note
+This expression assumes no biases.
+Main arguments:
+network: Linear network defined as a list of Equinox Linear layers.x: Network input.y: Network output.Returns:
+List of theoretical activities for each layer.
+ +jpc.pc_energy_fn(params: Tuple[PyTree[Callable], Optional[PyTree[Callable]]], activities: PyTree[ArrayLike], y: ArrayLike, x: Optional[ArrayLike] = None, loss: str = 'MSE', record_layers: bool = False) -> Array | Array
+
+
+¤Computes the free energy for a feedforward neural network of the form
+given parameters \(θ\), free activities \(\mathbf{z}\), output +\(\mathbf{z}_L = \mathbf{y}\) and optional input \(\mathbf{z}_0 = \mathbf{x}\) +for supervised training. The activity of each layer \(\mathbf{z}_\ell\) is +some function of the previous layer, e.g. +ReLU\((W_\ell \mathbf{z}_{\ell-1} + \mathbf{b}_\ell)\) for a fully connected +layer with biases and ReLU as activation.
+Note
+The input \(x\) and output \(y\) correspond to the prior and observation of +the generative model, respectively.
+Main arguments:
+params: Tuple with callable model (e.g. neural network) layers and
+ optional skip connections.activities: List of activities for each layer free to vary.y: Observation or target of the generative model.x: Optional prior of the generative model (for supervised training).Other arguments:
+loss: Loss function to use at the output layer (mean squared error
+ 'MSE' vs cross-entropy 'CE').
+ ??? cite "Reference"
@article{tscshantz2023hybrid,
+ title={Hybrid predictive coding: Inferring, fast and slow},
+ author={Tscshantz, Alexander and Millidge, Beren and Seth, Anil K and Buckley, Christopher L},
+ journal={PLoS Computational Biology},
+ volume={19},
+ number={8},
+ pages={e1011280},
+ year={2023},
+ publisher={Public Library of Science San Francisco, CA USA}
+}
+record_layers: If True, returns energies for each layer.
+Returns:
+The total or layer-wise energy normalised by the batch size.
+ +jpc.hpc_energy_fn(model: PyTree[typing.Callable], equilib_activities: PyTree[ArrayLike], amort_activities: PyTree[ArrayLike], x: ArrayLike, y: Optional[ArrayLike] = None, record_layers: bool = False) -> Array | Array
+
+
+¤Computes the free energy of an amortised PC network
+given the equilibrated activities of the generator \(\mathbf{z}^*\) (target +for the amortiser), the feedforward guesses of the amortiser +\(\hat{\mathbf{z}}\), the amortiser's parameters \(θ\), input +\(\mathbf{z}_0 = \mathbf{x}\), and optional output +\(\mathbf{z}_L = \mathbf{y}\) for supervised training.
+Note
+The input \(x\) and output \(y\) are reversed compared to pc_energy_fn
+(\(x\) is the generator's target and \(y\) is its optional input or prior).
+Just think of \(x\) and \(y\) as the actual input and output of the
+amortiser, respectively.
Main arguments:
+model: List of callable model (e.g. neural network) layers.equilib_activities: List of equilibrated activities reached by the
+ generator and target for the amortiser.amort_activities: List of amortiser's feedforward guesses
+ (initialisation) for the network activities.x: Input to the amortiser.y: Optional target of the amortiser (for supervised training).Other arguments:
+record_layers: If True, returns energies for each layer.Returns:
+The total or layer-wise energy normalised by batch size.
+ +Note
+There are two similar functions to compute the activity gradient:
+jpc.neg_activity_grad and jpc.compute_activity_grad. The first is used
+by jpc.solve_inference as gradient flow, while the second is for
+compatibility with discrete optax optimisers such as gradient descent.
jpc.neg_activity_grad(t: float | int, activities: PyTree[ArrayLike], args: Tuple[Tuple[PyTree[Callable], Optional[PyTree[Callable]]], ArrayLike, Optional[ArrayLike], str, diffrax._step_size_controller.base.AbstractStepSizeController]) -> PyTree[Array]
+
+
+¤Computes the negative gradient of the energy with respect to the activities \(- \partial \mathcal{F} / \partial \mathbf{z}\).
+This defines an ODE system to be integrated by solve_pc_inference.
Main arguments:
+t: Time step of the ODE system, used for downstream integration by
+ diffrax.diffeqsolve.activities: List of activities for each layer free to vary.args: 5-Tuple with
+ (i) Tuple with callable model layers and optional skip connections,
+ (ii) network output (observation),
+ (iii) network input (prior),
+ (iv) Loss specified at the output layer (MSE vs cross-entropy), and
+ (v) diffrax controller for step size integration.Returns:
+List of negative gradients of the energy w.r.t. the activities.
+ +jpc.compute_activity_grad(params: Tuple[PyTree[Callable], Optional[PyTree[Callable]]], activities: PyTree[ArrayLike], y: ArrayLike, x: Optional[ArrayLike], loss_id: str = 'MSE') -> PyTree[Array]
+
+
+¤Computes the gradient of the energy with respect to the activities \(\partial \mathcal{F} / \partial \mathbf{z}\).
+Note
+This function differs from neg_activity_grad, which computes the
+negative gradients, and is called in update_activities for use of
+any Optax optimiser.
Main arguments:
+params: Tuple with callable model layers and optional skip connections.activities: List of activities for each layer free to vary.y: Observation or target of the generative model.x: Optional prior of the generative model.Other arguments:
+loss_id: Loss function for the output layer (mean squared error 'MSE'
+ vs cross-entropy 'CE').energy_fn: Free energy to take the gradient of.Returns:
+List of negative gradients of the energy w.r.t. the activities.
+ +jpc.compute_pc_param_grads(params: Tuple[PyTree[Callable], Optional[PyTree[Callable]]], activities: PyTree[ArrayLike], y: ArrayLike, x: Optional[ArrayLike] = None, loss_id: str = 'MSE') -> Tuple[PyTree[Array], PyTree[Array]]
+
+
+¤Computes the gradient of the PC energy with respect to model parameters \(\partial \mathcal{F} / \partial θ\).
+Main arguments:
+params: Tuple with callable model layers and optional skip connections.activities: List of activities for each layer free to vary.y: Observation or target of the generative model.x: Optional prior of the generative model.Other arguments:
+loss_id: Loss function for the output layer (mean squared error 'MSE'
+ vs cross-entropy 'CE').Returns:
+List of parameter gradients for each network layer.
+ +jpc.compute_hpc_param_grads(model: PyTree[typing.Callable], equilib_activities: PyTree[ArrayLike], amort_activities: PyTree[ArrayLike], x: ArrayLike, y: Optional[ArrayLike] = None) -> PyTree[Array]
+
+
+¤Computes the gradient of the hybrid energy with respect to an amortiser's parameters \(\partial \mathcal{F} / \partial θ\).
+Main arguments:
+model: List of callable model (e.g. neural network) layers.equilib_activities: List of equilibrated activities reached by the
+ generator and target for the amortiser.amort_activities: List of amortiser's feedforward guesses
+ (initialisation) for the network activities.x: Input to the amortiser.y: Optional target of the amortiser (for supervised training).Note
+The input \(x\) and output \(y\) are reversed compared to compute_pc_param_grads
+(\(x\) is the generator's target and \(y\) is its optional input or prior).
+Just think of \(x\) and \(y\) as the actual input and output of the
+amortiser, respectively.
Returns:
+List of parameter gradients for each network layer.
+ +jpc.solve_inference(params: Tuple[PyTree[Callable], Optional[PyTree[Callable]]], activities: PyTree[ArrayLike], output: ArrayLike, input: Optional[ArrayLike] = None, loss_id: str = 'MSE', solver: AbstractSolver = Heun(scan_kind=None), max_t1: int = 20, dt: float | int = None, stepsize_controller: AbstractStepSizeController = PIDController(rtol=0.001,atol=0.001,pcoeff=0,icoeff=1,dcoeff=0,dtmin=None,dtmax=None,force_dtmin=True,step_ts=None,jump_ts=None,factormin=0.2,factormax=10.0,norm=<function rms_norm>,safety=0.9,error_order=None), record_iters: bool = False, record_every: int = None) -> PyTree[Array]
+
+
+¤Solves the inference (activity) dynamics of a predictive coding network.
+This is a wrapper around diffrax.diffeqsolve to integrate the gradient
+ODE system _neg_activity_grad defining the PC inference dynamics
where \(\mathcal{F}\) is the free energy, \(\mathbf{z}\) are the activities, +with \(\mathbf{z}_L\) clamped to some target and \(\mathbf{z}_0\) optionally +equal to some prior.
+Main arguments:
+params: Tuple with callable model layers and optional skip connections.activities: List of activities for each layer free to vary.output: Observation or target of the generative model.input: Optional prior of the generative model.Other arguments:
+loss: Loss function to use at the output layer (mean squared error
+ 'MSE' vs cross-entropy 'CE').solver: Diffrax (ODE) solver to be used. Default is Heun, a 2nd order
+ explicit Runge--Kutta method.max_t1: Maximum end of integration region (500 by default).dt: Integration step size. Defaults to None since the default
+ stepsize_controller will automatically determine it.stepsize_controller: diffrax controller for step size integration.
+ Defaults to PIDController. Note that the relative and absolute
+ tolerances of the controller will also determine the steady state to
+ terminate the solver.record_iters: If True, returns all integration steps.record_every: int determining the sampling frequency the integration
+ steps.Returns:
+List with solution of the activity dynamics for each layer.
+ +Info
+JPC provides 3 standard ways of initialising the activities: a feedforward +pass, randomly, or using an amortised network.
+jpc.init_activities_with_ffwd(model: PyTree[typing.Callable], input: ArrayLike, skip_model: Optional[PyTree[Callable]] = None) -> PyTree[Array]
+
+
+¤Initialises layers' activity with a feedforward pass +\(\{ f_\ell(\mathbf{z}_{\ell-1}) \}_{\ell=1}^L\) where \(\mathbf{z}_0 = \mathbf{x}\) is +the input.
+Main arguments:
+model: List of callable model (e.g. neural network) layers.input: input to the model.Other arguments:
+skip_model: Optional skip connection model.Returns:
+List with activity values of each layer.
+ +jpc.init_activities_from_normal(key: PRNGKeyArray, layer_sizes: PyTree[int], mode: str, batch_size: int, sigma: Array = 0.05) -> PyTree[Array]
+
+
+¤Initialises network activities from a zero-mean Gaussian \(\sim \mathcal{N}(0, \sigma^2)\).
+Main arguments:
+key: jax.random.PRNGKey for sampling.layer_sizes: List with dimension of all layers (input, hidden and
+ output).mode: If supervised, all hidden layers are initialised. If
+ unsupervised the input layer \(\mathbf{z}_0\) is also initialised.batch_size: Dimension of data batch.sigma: Standard deviation for Gaussian to sample activities from.
+ Defaults to 5e-2.Returns:
+List of randomly initialised activities for each layer.
+ +jpc.init_activities_with_amort(amortiser: PyTree[typing.Callable], generator: PyTree[typing.Callable], input: ArrayLike) -> PyTree[Array]
+
+
+¤Initialises layers' activity with an amortised network +\(\{ f_{L-\ell+1}(\mathbf{z}_{L-\ell}) \}_{\ell=1}^L\) where \(\mathbf{z}_0 = \mathbf{y}\) is +the input or generator's target.
+Note
+The output order is reversed for downstream use by the generator.
+Main arguments:
+amortiser: List of callable layers for model amortising the inference
+ of the generator.generator: List of callable layers for the generative model.input: Input to the amortiser.Returns:
+List with amortised initialisation of each layer.
+ +jpc.test_discriminative_pc(model: PyTree[typing.Callable], output: ArrayLike, input: ArrayLike, loss: str = 'MSE', skip_model: Optional[PyTree[Callable]] = None) -> Tuple[Array, Array]
+
+
+¤Computes test metrics for a discriminative predictive coding network.
+Main arguments:
+model: List of callable model (e.g. neural network) layers.output: Observation or target of the generative model.input: Optional prior of the generative model.Other arguments:
+loss: - loss: Loss function to use at the output layer (mean
+ squared error 'MSE' vs cross-entropy 'CE').skip_model: Optional list of callable skip connection functions.Returns:
+Test loss and accuracy of output predictions.
+ +jpc.test_generative_pc(model: PyTree[typing.Callable], output: ArrayLike, input: ArrayLike, key: PRNGKeyArray, layer_sizes: PyTree[int], batch_size: int, sigma: Array = 0.05, ode_solver: AbstractSolver = Heun(scan_kind=None), max_t1: int = 500, dt: Array | int = None, stepsize_controller: AbstractStepSizeController = PIDController(rtol=0.001,atol=0.001,pcoeff=0,icoeff=1,dcoeff=0,dtmin=None,dtmax=None,force_dtmin=True,step_ts=None,jump_ts=None,factormin=0.2,factormax=10.0,norm=<function rms_norm>,safety=0.9,error_order=None), skip_model: Optional[PyTree[Callable]] = None) -> Tuple[Array, Array]
+
+
+¤Computes test metrics for a generative predictive coding network.
+Gets output predictions (e.g. of an image given a label) with a feedforward +pass and calculates accuracy of inferred input (e.g. of a label given an +image).
+Main arguments:
+model: List of callable model (e.g. neural network) layers.output: Observation or target of the generative model.input: Optional prior of the generative model.key: jax.random.PRNGKey for random initialisation of activities.layer_sizes: Dimension of all layers (input, hidden and output).batch_size: Dimension of data batch for activity initialisation.Other arguments:
+sigma: Standard deviation for Gaussian to sample activities from.
+ Defaults to 5e-2.ode_solver: Diffrax ODE solver to be used. Default is Heun, a 2nd order
+ explicit Runge--Kutta method.max_t1: Maximum end of integration region (500 by default).dt: Integration step size. Defaults to None since the default
+ stepsize_controller will automatically determine it.stepsize_controller: diffrax controller for step size integration.
+ Defaults to PIDController. Note that the relative and absolute
+ tolerances of the controller will also determine the steady state to
+ terminate the solver.Returns:
+Accuracy and output predictions.
+ +jpc.test_hpc(generator: PyTree[typing.Callable], amortiser: PyTree[typing.Callable], output: ArrayLike, input: ArrayLike, key: PRNGKeyArray, layer_sizes: PyTree[int], batch_size: int, sigma: Array = 0.05, ode_solver: AbstractSolver = Heun(scan_kind=None), max_t1: int = 500, dt: Array | int = None, stepsize_controller: AbstractStepSizeController = PIDController(rtol=0.001,atol=0.001,pcoeff=0,icoeff=1,dcoeff=0,dtmin=None,dtmax=None,force_dtmin=True,step_ts=None,jump_ts=None,factormin=0.2,factormax=10.0,norm=<function rms_norm>,safety=0.9,error_order=None)) -> Tuple[Array, Array, Array, Array]
+
+
+¤Computes test metrics for hybrid predictive coding trained in a supervised manner.
+Calculates input accuracy of (i) amortiser, (ii) generator, and (iii) +hybrid (amortiser + generator). Also returns output predictions (e.g. of +an image given a label) with a feedforward pass of the generator.
+Note
+The input and output of the generator are the output and input of the +amortiser, respectively.
+Main arguments:
+generator: List of callable layers for the generative model.amortiser: List of callable layers for model amortising the inference
+ of the generator.output: Observation or target of the generative model.input: Optional prior of the generator, target for the amortiser.key: jax.random.PRNGKey for random initialisation of activities.layer_sizes: Dimension of all layers (input, hidden and output).batch_size: Dimension of data batch for initialisation of activities.Other arguments:
+sigma: Standard deviation for Gaussian to sample activities from.
+ Defaults to 5e-2.ode_solver: Diffrax ODE solver to be used. Default is Heun, a 2nd order
+ explicit Runge--Kutta method.max_t1: Maximum end of integration region (500 by default).dt: Integration step size. Defaults to None since the default
+ stepsize_controller will automatically determine it.stepsize_controller: diffrax controller for step size integration.
+ Defaults to PIDController. Note that the relative and absolute
+ tolerances of the controller will also determine the steady state to
+ terminate the solver.Returns:
+Accuracies of all models and output predictions.
+ +jpc.make_pc_step(model: PyTree[typing.Callable], optim: optax._src.base.GradientTransformation | optax._src.base.GradientTransformationExtraArgs, opt_state: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, Iterable[ArrayTree], Mapping[Any, ArrayTree]], output: ArrayLike, input: Optional[ArrayLike] = None, loss_id: str = 'MSE', ode_solver: AbstractSolver = Heun(scan_kind=None), max_t1: int = 20, dt: Array | int = None, stepsize_controller: AbstractStepSizeController = PIDController(rtol=0.001,atol=0.001,pcoeff=0,icoeff=1,dcoeff=0,dtmin=None,dtmax=None,force_dtmin=True,step_ts=None,jump_ts=None,factormin=0.2,factormax=10.0,norm=<function rms_norm>,safety=0.9,error_order=None), skip_model: Optional[PyTree[Callable]] = None, key: Optional[PRNGKeyArray] = None, layer_sizes: Optional[PyTree[int]] = None, batch_size: Optional[int] = None, sigma: Array = 0.05, record_activities: bool = False, record_energies: bool = False, record_every: int = None, activity_norms: bool = False, param_norms: bool = False, grad_norms: bool = False, calculate_accuracy: bool = False) -> Dict
+
+
+¤Updates network parameters with predictive coding.
+Main arguments:
+model: List of callable model (e.g. neural network) layers.optim: Optax optimiser, e.g. optax.sgd().opt_state: State of Optax optimiser.output: Observation or target of the generative model.input: Optional prior of the generative model.Note
+key, layer_sizes and batch_size must be passed if input is
+None, since unsupervised training will be assumed and activities need
+to be initialised randomly.
Other arguments:
+loss_id: Loss function for the output layer (mean squared error 'MSE'
+ vs cross-entropy 'CE').ode_solver: Diffrax ODE solver to be used. Default is Heun, a 2nd order
+ explicit Runge--Kutta method.max_t1: Maximum end of integration region (20 by default).dt: Integration step size. Defaults to None since the default
+ stepsize_controller will automatically determine it.stepsize_controller: diffrax controller for step size integration.
+ Defaults to PIDController. Note that the relative and absolute
+ tolerances of the controller will also determine the steady state to
+ terminate the solver.skip_model: Optional list of callable skip connection functions.key: jax.random.PRNGKey for random initialisation of activities.layer_sizes: Dimension of all layers (input, hidden and output).batch_size: Dimension of data batch for activity initialisation.sigma: Standard deviation for Gaussian to sample activities from for
+ random initialisation. Defaults to 5e-2.record_activities: If True, returns activities at every inference
+ iteration.record_energies: If True, returns layer-wise energies at every
+ inference iteration.record_every: int determining the sampling frequency the integration
+ steps.activity_norms: If True, computes l2 norm of the activities.param_norms: If True, computes l2 norm of the parameters.grad_norms: If True, computes l2 norm of parameter gradients.calculate_accuracy: If True, computes the training accuracy.Returns:
+Dict including model (and optional skip model) with updated parameters, +optimiser, updated optimiser state, loss, energies, activities, +and optionally other metrics (see other args above).
+Raises:
+ValueError for inconsistent inputs and invalid losses.jpc.make_hpc_step(generator: PyTree[typing.Callable], amortiser: PyTree[typing.Callable], optims: Tuple[optax._src.base.GradientTransformationExtraArgs], opt_states: Tuple[Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, Iterable[ArrayTree], Mapping[Any, ArrayTree]]], output: ArrayLike, input: Optional[ArrayLike] = None, ode_solver: AbstractSolver = Heun(scan_kind=None), max_t1: int = 300, dt: Array | int = None, stepsize_controller: AbstractStepSizeController = PIDController(rtol=0.001,atol=0.001,pcoeff=0,icoeff=1,dcoeff=0,dtmin=None,dtmax=None,force_dtmin=True,step_ts=None,jump_ts=None,factormin=0.2,factormax=10.0,norm=<function rms_norm>,safety=0.9,error_order=None), record_activities: bool = False, record_energies: bool = False) -> Dict
+
+
+¤Updates parameters of a hybrid predictive coding network.
+@article{tscshantz2023hybrid,
+ title={Hybrid predictive coding: Inferring, fast and slow},
+ author={Tscshantz, Alexander and Millidge, Beren and Seth, Anil K and Buckley, Christopher L},
+ journal={PLoS Computational Biology},
+ volume={19},
+ number={8},
+ pages={e1011280},
+ year={2023},
+ publisher={Public Library of Science San Francisco, CA USA}
+}
+Note
+The input and output of the generator are the output and input of the +amortiser, respectively.
+Main arguments:
+generator: List of callable layers for the generative model.amortiser: List of callable layers for model amortising the inference
+ of the generator.optims: Optax optimisers (e.g. optax.sgd()), one for each model.opt_states: State of Optax optimisers, one for each model.output: Observation of the generator, input to the amortiser.input: Optional prior of the generator, target for the amortiser.Other arguments:
+ode_solver: Diffrax ODE solver to be used. Default is Heun, a 2nd order
+ explicit Runge--Kutta method..max_t1: Maximum end of integration region (300 by default).dt: Integration step size. Defaults to None since the default
+ stepsize_controller will automatically determine it.stepsize_controller: diffrax controller for step size integration.
+ Defaults to PIDController. Note that the relative and absolute
+ tolerances of the controller will also determine the steady state to
+ terminate the solver.record_activities: If True, returns activities at every inference
+ iteration.record_energies: If True, returns layer-wise energies at every
+ inference iteration.Returns:
+Dict including models with updated parameters, optimiser and state for each +model, model activities, last inference step for the generator, MSE losses, +and energies.
+ +jpc.update_activities(params: Tuple[PyTree[Callable], Optional[PyTree[Callable]]], activities: PyTree[ArrayLike], optim: optax._src.base.GradientTransformation | optax._src.base.GradientTransformationExtraArgs, opt_state: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, Iterable[ArrayTree], Mapping[Any, ArrayTree]], output: ArrayLike, input: Optional[ArrayLike] = None) -> Dict
+
+
+¤Updates activities of a predictive coding network with a given Optax optimiser.
+Main arguments:
+params: Tuple with callable model layers and optional skip connections.activities: List of activities for each layer free to vary.optim: Optax optimiser, e.g. optax.sgd().opt_state: State of Optax optimiser.output: Observation or target of the generative model.input: Optional prior of the generative model.Returns:
+Dictionary with energy, updated activities, activity gradients, optimiser, +and updated optimiser state.
+ +jpc.update_params(params: Tuple[PyTree[Callable], Optional[PyTree[Callable]]], activities: PyTree[ArrayLike], optim: optax._src.base.GradientTransformation | optax._src.base.GradientTransformationExtraArgs, opt_state: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, Iterable[ArrayTree], Mapping[Any, ArrayTree]], output: ArrayLike, input: Optional[ArrayLike] = None) -> Dict
+
+
+¤Updates parameters of a predictive coding network with a given Optax optimiser.
+Main arguments:
+params: Tuple with callable model layers and optional skip connections.activities: List of activities for each layer free to vary.optim: Optax optimiser, e.g. optax.sgd().opt_state: State of Optax optimiser.output: Observation or target of the generative model.input: Optional prior of the generative model.Returns:
+Dictionary with model (and optional skip model) with updated parameters, +parameter gradients, optimiser, and updated optimiser state.
+ +jpc.make_mlp(key: PRNGKeyArray, layer_sizes: PyTree[int], act_fn: str, use_bias: bool = True) -> PyTree[typing.Callable]
+
+
+¤Creates a multi-layer perceptron compatible with predictive coding updates.
+Main arguments:
+key: jax.random.PRNGKey for parameter initialisation.layer_sizes: Dimension of all layers (input, hidden and output).
+ Options are linear, tanh and relu.act_fn: Activation function for all layers except the output.use_bias: True by default.Returns:
+List of callable fully connected layers.
+ +jpc.get_act_fn(name: str) -> Callable
+
+
+¤jpc.mse_loss(preds: ArrayLike, labels: ArrayLike) -> Array
+
+
+¤jpc.cross_entropy_loss(logits: ArrayLike, labels: ArrayLike) -> Array
+
+
+¤jpc.compute_accuracy(truths: ArrayLike, preds: ArrayLike) -> Array
+
+
+¤jpc.get_t_max(activities_iters: PyTree[Array]) -> Array
+
+
+¤jpc.compute_infer_energies(params: Tuple[PyTree[Callable], Optional[PyTree[Callable]]], activities_iters: PyTree[Array], t_max: Array, y: ArrayLike, x: Optional[ArrayLike] = None, loss: str = 'MSE') -> PyTree[Array]
+
+
+¤Calculates layer energies during predictive coding inference.
+Main arguments:
+params: Tuple with callable model layers and optional skip connections.activities_iters: Layer-wise activities at every inference iteration.
+ Note that each set of activities will have 4096 steps as first
+ dimension by diffrax default.t_max: Maximum number of inference iterations to compute energies for.y: Observation or target of the generative model.x: Optional prior of the generative model.Other arguments:
+loss: Loss function specified at the output layer (mean squared error
+ 'MSE' vs cross-entropy 'CE').Returns:
+List of layer-wise energies at every inference iteration.
+ +jpc.compute_activity_norms(activities: PyTree[Array]) -> Array
+
+
+¤Calculates l2 norm of activities at each layer.
+ +jpc.compute_param_norms(params)
+
+
+¤Calculates l2 norm of all model parameters.
+ +