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]#

Initialize tn4ml.models.Model.

Return type:

None

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

Save 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]#

Return number of parameters of the model.

Return type:

int

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

Configure model for training with specific parameters.

Parameters:
  • kwargs (dict)

  • keys (Configuration parameters. Supported)

  • strategy (-)

  • ('global' (Training strategy)

  • 'sweeps'

  • 'local'

  • 'dmrg'

  • 'dmrg-like')

  • optimizer (-) – Optimization algorithm to use.

  • loss (-) – Loss function for training.

  • train_type (-) – Type of training (from tn4ml.util.TrainingType)

  • learning_rate (-) – Learning rate for optimizer.

  • gradient_transforms (-) – Sequence of gradient transformations for optax

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

Return type:

None

Examples

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

Predict 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=None, 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=None, batch_size=64, shuffle=False, normalize=False, accuracy_fn=None, dtype=<class 'jax.numpy.float64'>, seed=42, alternate_flip=False)[source][source]#

Calculate 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.

  • accuracy_fn (Callable) – Function applied to raw model outputs before class labels are extracted. If it returns class scores, the predicted class is selected with argmax; if it returns class labels, those labels are compared directly.

  • 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]#

Update 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]#

Compute 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]#

Compute 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]#

Create functions for training steps and gradients.

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=None, 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, accuracy_fn=None, dtype=<class 'jax.numpy.float64'>, shuffle=False, seed=42, alternate_flip=False)[source][source]#

Perform 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.

  • accuracy_fn (Callable) – Function applied to raw model outputs before validation accuracy labels are extracted. Passed to 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=None, evaluate_type=TrainingType.UNSUPERVISED, return_list=False, metric=None, dtype=<class 'jax.numpy.float64'>, shuffle=False, seed=42, alternate_flip=False)[source][source]#

Evaluate 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]#

Convert tensor network to pytree structure.

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]#

Load 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=None, 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]#

Normalize 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]#

Calculate 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#

Get lower indices for output sites.

get_orders()[source][source]#

Get tensor index orders.

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]#

Apply this SMPO to another SMPO or MPS.

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]#

Initialize 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 tensors of tn4ml.models.mps.MatrixProductState.

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]#

Initialize tn4ml.models.Model.

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

Normalize 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]#

Initialize 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]#

Copy the model.

Return type:

Model of the same type.

Parameters:
  • virtual (bool)

  • deep (bool)

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

Calculate 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]#

Normalize 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