Model

Model#

class tn4ml.models.Model[source][source]#

Bases: TensorNetwork

tn4ml.models.Model class models training model of class quimb.tensor.tensor_core.TensorNetwork.

loss#

Loss function. See tn4ml.metrics for examples.

Type:

Callable, or None

strategy#

Strategy for computing gradients.

Type:

tn4ml.strategy.Strategy

optimizer#

Type of optimizer matching names of optimizers from optax.

Type:

str

learning_rate#

Learning rate for optimizer.

Type:

float

train_type#

Type of training: 0 = ‘unsupervised’ or 1 =’supervised’, 2 = ‘target TN’.

Type:

int

gradient_transforms#

Sequence of gradient transformations.

Type:

sequence

device#

Device for computation, e.g. (‘cpu’, 0) or (‘gpu’, 0). [0] = device name, [1] = device index.

Type:

(str, int)

opt_state#

State of optimizer.

Type:

Any

__init__()[source][source]#

Constructor method for tn4ml.models.Model class.

save(model_name, dir_name='~', tn=False)[source][source]#

Saves tn4ml.models.Model to pickle file.

Parameters:

  • model_name (str) – Name of Model.

  • dir_name (str) – Directory for saving Model.

  • tn (bool) – If True, model object is TensorNetwork.

nparams()[source][source]#

Returns number of parameters of the model.

Return type:

int

configure(**kwargs)[source][source]#

Configures model for training with specific parameters.

Parameters:

kwargs (dict) –

Configuration parameters.

Including: - strategy: str or Strategy object Training strategy (‘global’, ‘sweeps’, ‘local’, ‘dmrg’, ‘dmrg-like’) - optimizer: callable or optax optimizer

Optimization algorithm to use

  • loss: callable

    Loss function for training

  • train_type: int

    Type of training (from tn4ml.util.TrainingType)

  • learning_rate: float

    Learning rate for optimizer

  • gradient_transforms: sequence

    Sequence of gradient transformations for optax

  • device: (str, int)

    Device for computation, e.g. (‘cpu’, 0) or (‘gpu’, 0). [0] = device name, [1] = device index.

Examples

>>> model.configure(strategy='global', optimizer=optax.adam, learning_rate=0.01, loss=tn4ml.metrics.LogQuadNorm, train_type=TrainingType.UNSUPERVISED)
predict(sample, embedding=<tn4ml.embeddings.TrigonometricEmbedding object>, return_tn=False, normalize=False)[source][source]#

Predicts the output of the model.

Parameters:

  • sample (numpy.ndarray) – Input data.

  • embedding (tn4ml.embeddings.Embedding) – Data embedding function.

  • return_tn (bool) – If True, returns tensor network, otherwise returns data. Useful when you want to vmap over predict function.

  • normalize (bool)

Returns:

Output of the model.

Return type:

quimb.tensor.tensor_core.TensorNetwork

forward(data, embedding=<tn4ml.embeddings.TrigonometricEmbedding object>, batch_size=64, normalize=False, dtype=<class 'jax.numpy.float64'>, seed=42, alternate_flip=False)[source][source]#

Forward pass of the model.

Parameters:

  • data (jax.numpy.ndarray) – Input data.

  • y_true (jax.numpy.ndarray) – Target class vector.

  • embedding (tn4ml.embeddings.Embedding) – Data embedding function.

  • batch_size (int) – Batch size for data processing.

  • normalize (bool) – If True, the model output is normalized.

  • dtype (Any) – Data type of input data.

  • seed (int) – Random seed for data shuffling.

  • alternate_flip (bool)

Returns:

Output of the model.

Return type:

jax.numpy.ndarray

accuracy(data, y_true=None, embedding=<tn4ml.embeddings.TrigonometricEmbedding object>, batch_size=64, shuffle=False, normalize=False, dtype=<class 'jax.numpy.float64'>, seed=42, alternate_flip=False)[source][source]#

Calculates accuracy for supervised learning.

Parameters:

  • model (tn4ml.models.Model) – Tensor Network model.

  • data (numpy.ndarray) – Input data.

  • y_true (numpy.ndarray) – Target class vector.

  • embedding (tn4ml.embeddings.Embedding) – Data embedding function.

  • batch_size (int) – Batch size for data processing.

  • normalize (bool) – If True, the model output is normalized in predict function.

  • dtype (Any) – Data type of input data.

  • seed (int) – Random seed for data shuffling.

  • shuffle (bool)

  • alternate_flip (bool)

Return type:

float

update_tensors(params)[source][source]#

Updates tensors of the model with new parameters.

Parameters:

  • params (sequence of jax.numpy.ndarray) – New parameters of the model.

  • sitetags (sequence of str, or default None) – Names of tensors for differentiation (for Sweeping strategy).

Return type:

None

compute_entropy(data, embedding)[source][source]#

Computes entropy of the model.

Parameters:

  • data (jax.numpy.ndarray) – Input data.

  • embedding (tn4ml.embeddings.Embedding) – Data embedding function.

Returns:

Entropy of the model.

Return type:

float

compute_entropy_batch(data, embedding)[source][source]#

Computes entropy of the model for a batch of data.

Parameters:

  • data (jax.numpy.ndarray) – Input data.

  • embedding (tn4ml.embeddings.Embedding) – Data embedding function.

Returns:

Entropy of the model.

Return type:

float

create_train_step(params, loss_func)[source][source]#

Creates function for calculating value and gradients of loss, and function for one step in training procedure. Initializes the optimizer and creates optimizer state.

Parameters:

  • params (sequence of jax.numpy.ndarray) – Parameters of the model.

  • loss_func (function) – Loss function.

  • grads_func (function) – Function for calculating gradients of loss.

Returns:

  • train_step (function) – Function to perform one training step.

  • opt_state (tuple) – State of optimizer at the initialization.

train(inputs=None, val_inputs=None, targets=None, val_targets=None, tn_target=None, batch_size=None, epochs=1, embedding=<tn4ml.embeddings.TrigonometricEmbedding object>, normalize=False, canonize=(False, None), time_limit=None, earlystop=None, gradient_clip_threshold=None, val_batch_size=None, eval_metric=None, display_val_acc=False, dtype=<class 'jax.numpy.float64'>, shuffle=False, seed=42, alternate_flip=False)[source][source]#

Performs the training procedure of tn4ml.models.Model.

Parameters:

  • inputs (sequence of numpy.ndarray) – Data used for training procedure.

  • val_inputs (sequence of numpy.ndarray) – Data used for validation.

  • targets (sequence of numpy.ndarray) – Targets for training procedure (if training is supervised).

  • val_targets (sequence of numpy.ndarray) – Targets for validation (if training is supervised).

  • tn_target (quimb.tensor.tensor_core.TensorNetwork or any specialized TN class from quimb.tensor module) – Target tensor network for training.

  • batch_size (int, or default None) – Number of samples per gradient update.

  • epochs (int) – Number of epochs for training.

  • embedding (tn4ml.embeddings.Embedding) – Data embedding function.

  • normalize (bool) – If True, the model is normalized after each iteration.

  • canonize (tuple([bool, int])) – tuple indicating is model canonized after each iteration. Example: (True, 0) - model is canonized in canonization center = 0.

  • time_limit (int) – Time limit on model’s training in seconds.

  • earlystop (tn4ml.util.EarlyStopping) – Early stopping training when monitored metric stopped improving.

  • gradient_clip_threshold (float) – Threshold for gradient clipping.

  • val_batch_size (int) – Number of samples per validation batch.

  • display_val_acc (bool) – If True, displays validation accuracy.

  • alternate_flip (bool) – If True, flips every other batch along axis=1.

  • eval_metric (Callable | None)

  • dtype (Any)

  • shuffle (bool | None)

  • seed (int | None)

Returns:

history – Records training loss and metric values.

Return type:

dict

evaluate(inputs=None, targets=None, tn_target=None, batch_size=None, embedding=<tn4ml.embeddings.TrigonometricEmbedding object>, evaluate_type=TrainingType.UNSUPERVISED, return_list=False, metric=None, dtype=<class 'jax.numpy.float64'>, shuffle=False, seed=42, alternate_flip=False)[source][source]#

Evaluates the model on the data.

Parameters:

  • inputs (sequence of numpy.ndarray) – Data used for evaluation.

  • targets (sequence of numpy.ndarray) – Targets for evaluation (if evaluation is supervised).

  • tn_target (quimb.tensor.tensor_core.TensorNetwork or any specialized TN class from quimb) – Target tensor network for evaluation.

  • batch_size (int, or default None) – Number of samples per evaluation.

  • embedding (tn4ml.embeddings.Embedding) – Data embedding function.

  • evaluate_type (int) – Type of evaluation: 0 = ‘unsupervised’ or 1 =’unsupervised’.

  • return_list (bool) – If True, returns list of loss values for each batch.

  • metric (function) – Metric function for evaluation.

  • dtype (Any) – Data type of input data.

  • shuffle (bool) – If True, data is shuffled.

  • seed (int) – Random seed for data shuffling.

  • alternate_flip (bool | None)

Returns:

Loss value.

Return type:

float

convert_to_pytree()[source][source]#

Converts tensor network to pytree structure and returns its skeleon. Reference to quimb.tensor.pack().

Returns:

  • pytree (dict)

  • skeleton (Tensor, TensorNetwork, or similar) - A copy of obj with all references to the original data removed.

tn4ml.models.load_model(model_name, dir_name=None)[source][source]#

Loads the Model from pickle file.

Parameters:

  • model_name (str) – Name of the model.

  • dir_name (str) – Directory where model is stored.

Return type:

tn4ml.models.Model or subclass

class tn4ml.models.smpo.SpacedMatrixProductOperator[source][source]#

Bases: TensorNetwork1DOperator, TensorNetwork1DFlat, Model

A MatrixProductOperator with a decimated number of output indices. See quimb.tensor.tensor_1d.MatrixProductOperator for explanation of other attributes and methods.

__init__(arrays, output_inds=[], shape='lrud', site_tag_id='I{}', tags=None, upper_ind_id='k{}', lower_ind_id='b{}', bond_name='bond{}', **tn_opts)[source][source]#

Create a MatrixProductOperator with a decimated number of output indices.

Return type:

None

arrays#

The arrays defining the operator.

Type:

tuple of array_like

output_inds#

Indexes of tensors which have output indices. From 0 to n. If spacing is not even.

Type:

array of int

shape#

The shape of the tensors, e.g. ‘lurd’ or ‘lrud’.

Type:

str

site_tag_id#

The format string for the site tags.

Type:

str

tags#

Global tags to add to all tensors.

Type:

str or sequence of str

upper_ind_id#

The format string for the upper virtual indices.

Type:

str

lower_ind_id#

The format string for the lower virtual indices.

Type:

str

bond_name#

The format string for the bond names.

Type:

str

tn_opts#

Supplied to quimb.tensor.tensor_core.TensorNetwork.

Type:

optional

normalize(insert=None, output_inds=None)[source][source]#

Function for normalizing tensors of tn4ml.models.smpo.SpacedMatrixProductOperator.

Parameters:

insert (int) – Index of tensor divided by norm. Default = None. When None the norm division is distributed across all tensors.

Return type:

None

norm(**contract_opts)[source][source]#

Calculates norm of tn4ml.models.smpo.SpacedMatrixProductOperator.

Parameters:

contract_opts (Optional) – Arguments passed to contract().

Returns:

Norm of tn4ml.models.smpo.SpacedMatrixProductOperator

Return type:

float

property spacing: int#

Spacing paramater, or space between output indices in number of sites.

property spacings: list#

Spacings paramater, or space between output indices in number of sites.

property lower_inds#

Return a tuple of all lower indices.

get_orders()[source][source]#
Return type:

list

apply_mps(tn_vec, normalize_on_contract=True, compress=False, **compress_opts)[source][source]#

Version of quimb.tensor.tensor_1d.MatrixProductOperator._apply_mps() for tn4ml.models.smpo.SpacedMatrixProductOperator.

Parameters:

  • tn_op (quimb.tensor.tensor_core.TensorNetwork) – The tensor network representing the operator.

  • tn_vec (quimb.tensor.tensor_core.TensorNetwork, or quimb.tensor.tensor_1d.MatrixProductState) – The tensor network representing the vector.

  • compress (bool) – Whether to compress the resulting tensor network.

  • compress_opts (optional) – Options to pass to tn_vec.compress.

Return type:

quimb.tensor.tensor_1d.MatrixProductState

apply_smpo(tn_op_2, trace=True, compress=False, **compress_opts)[source][source]#

Version of quimb.tensor.tensor_1d.MatrixProductOperator._apply_mpo() for tn4ml.models.smpo.SpacedMatrixProductOperator.

Parameters:
Return type:

quimb.tensor.tensor_1d.TensorNetwork1D

apply(other, normalize_on_contract=False, compress=False, **compress_opts)[source][source]#

Version of quimb.tensor.tensor_1d.MatrixProductOperator.apply() for tn4ml.models.smpo.SpacedMatrixProductOperator. Act with this SMPO on another SMPO or MPS, such that the resulting object has the same tensor network structure/indices as other. For an MPS:

For an SMPO:

The resulting TN will have the same structure/indices as other, but probably with larger bonds (depending on compression).

Parameters:

  • other (tn4ml.models.smpo.SpacedMatrixProductOperator, or quimb.tensor.tensor_1d.MatrixProductState.) – The object to act on.

  • compress (bool) – Whether to compress the resulting object.

  • compress_opts (optional) – Supplied to TensorNetwork1DFlat.compress().

Return type:

quimb.tensor.tensor_1d.MatrixProductOperator, or quimb.tensor.tensor_1d.MatrixProductState

class tn4ml.models.mps.MatrixProductState[source][source]#

Bases: Model, MatrixProductState

A Trainable MatrixProductState class. See quimb.tensor.tensor_1d.MatrixProductState for explanation of other attributes and methods.

__init__(arrays, **kwargs)[source][source]#

Initializes the MatrixProductState.

Parameters:

  • arrays (list of array_like) – The list of tensors, each of shape (D, D, d), where D is the bond dimension and d is the physical dimension.

  • **kwargs (dict) – Additional arguments to be passed to the parent class.

normalize(insert=None)[source][source]#

Normalize this MPS, optional with co-vector bra. For periodic MPS this uses transfer matrix SVD approximation with precision eps in order to be efficient. Inplace.

Parameters:

  • bra (MatrixProductState, optional) – If given, normalize this MPS with the same factor.

  • eps (float, optional) – If cyclic, precision to approximation transfer matrix with. Default: 1e-14.

  • insert (int, optional) – Insert the corrective normalization on this site, random if not given.

Returns:

old_norm – The old norm self.H @ self.

Return type:

float

class tn4ml.models.mpo.MatrixProductOperator[source][source]#

Bases: Model, MatrixProductOperator

A Trainable MatrixProductOperator class. See quimb.tensor.tensor_1d.MatrixProductOperator for explanation of other attributes and methods.

__init__(arrays, **kwargs)[source][source]#

Constructor method for tn4ml.models.Model class.

normalize(insert=None)[source][source]#

Function for normalizing tensors of tn4ml.models.mpo.MatrixProductOperator.

Parameters:

insert (int) – Index of tensor divided by norm. Default = None. When None the norm division is distributed across all tensors.

class tn4ml.models.tn.TensorNetwork[source][source]#

Bases: Model, TensorNetwork1DFlat

A Trainable TensorNetwork class. See quimb.tensor.tensor_core.TensorNetwork for explanation of other attributes and methods.

__init__(tensors, site_tag_id='I{}', cyclic=False, **kwargs)[source][source]#

Initializes tn4ml.models.tn.ParametrizedTensorNetwork.

Parameters:

  • tensors (list or TensorNetwork) – List of tensors of quimb.tensor.tensor_core.Tensor or :class:quimb.tensor.tensor_core.TensorNetwork.

  • kwargs (dict) – Additional arguments.

  • site_tag_id (str)

  • cyclic (bool)

canonize(where, cur_orthog='calc', info=None, bra=None, inplace=False)[source][source]#

Canonizes the tensor network.

copy(virtual=False, deep=False)[source][source]#

Copies the model.

Return type:

Model of the same type.

Parameters:

  • virtual (bool)

  • deep (bool)

norm(**contract_opts)[source][source]#

Calculates norm of tn4ml.models.tn.TensorNetwork.

Parameters:

contract_opts (Optional) – Arguments passed to contract().

Returns:

Norm of tn4ml.models.smpo.SpacedMatrixProductOperator

Return type:

float

normalize(insert=None)[source][source]#

Function for normalizing tensors of tn4ml.models.tn.TensorNetwork.

Parameters:

insert (int) – Index of tensor divided by norm. Default = None. When None the norm division is distributed across all tensors.

Return type:

None

Contents