AOPMM#

parametricmatrixmodels.modules.AOPMM

AOPMM#

add

Append a module to the end of the model.

add_module

Append a module to the end of the model.

append

Append a module to the end of the model.

append_module

Append a module to the end of the model.

astype

Convenience wrapper to set_precision using the dtype argument, returns self.

compile

Compile the model for training by compiling each module.

copy

Create a deep copy of the module.

deepcopy

Create a deep copy of the module.

deserialize

Deserialize the model from a dictionary.

freeze

Freeze the module parameters by setting trainable to False.

from_file

Load a model from a file and return an instance of the Model class.

get_hyperparameters

Get the hyperparameters of the model as a dictionary.

get_modules

Get the modules of the model.

get_num_trainable_floats

Returns the number of trainable floats in the module.

get_output_shape

Get the output shape of the model given an input shape.

get_params

Get the parameters of the model.

get_rng

get_state

Get the state of the model.

get_trainable_modules

Return a PyTree with the same structure as the model's modules, where each leaf node is a boolean indicating whether the corresponding module is trainable (i.e. mod.trainable == True and mod.get_params() is not None/()).

grad_input

Doc TODO

grad_params

Doc TODO

insert

Insert a module at a specific index in the model.

insert_module

Insert a module at a specific index in the model.

is_ready

Return True if the module is initialized and ready for training or inference.

load

Load the model from a file.

pop

Remove and return a module by key or index in the model.

pop_module_by_index

Remove and return a module at a specific index in the model.

pop_module_by_key

Remove and return a module by key or index in the model.

predict

Call the model with the input data.

prepend

Prepend a module to the beginning of the model.

prepend_module

Prepend a module to the beginning of the model.

reset

save

Save the model to a file.

save_compressed

Save the model to a compressed file.

serialize

Serialize the model to a dictionary.

set_hyperparameters

Set the hyperparameters of the model from a dictionary.

set_params

Set the parameters of the model from a PyTree of PyTrees of numpy arrays.

set_precision

Set the precision of the model parameters and states.

set_rng

Set the random key for the model.

set_state

Set the state of the model from a PyTree of PyTrees of numpy arrays.

train

unfreeze

Unfreeze the module parameters by setting trainable to True.

upgrade

Upgrade serialized module data to the current version.

name

Returns the name of the module, unless overridden, this is the class name.

trainable

Whether the module is trainable (i.e., whether its parameters should be updated during training).

alias of AffineObservablePMM