TensorBlock

class metatensor.torch.TensorBlock(values: Tensor, samples: Labels, components: List[Labels], properties: Labels)

Basic building block for a TensorMap.

A single block contains a n-dimensional torch.Tensor of values, and n sets of Labels (one for each dimension). The first dimension is the samples dimension, the last dimension is the properties dimension. Any intermediate dimension is called a component dimension.

Samples should be used to describe what we are representing, while properties should contain information about how we are representing it. Finally, components should be used to describe vectorial or tensorial components of the data.

A block can also contain gradients of the values with respect to a variety of parameters. In this case, each gradient is a TensorBlock with a separate set of samples and possibly components, but which shares the same property labels as the original TensorBlock.

See also

The pure Python version of this class metatensor.TensorBlock, and the differences between TorchScript and Python API for metatensor.

Parameters:

• values (Tensor) – tensor containing the values for this block

• samples (Labels) – labels describing the samples (first dimension of the array)

• components (List[Labels]) – list of labels describing the components (intermediate dimensions of the array). This should be an empty list for scalar/invariant data.

• properties (Labels) – labels describing the properties (last dimension of the array)

property values: Tensor

get the values for this block

property samples: Labels

Get the sample Labels for this block.

The entries in these labels describe the first dimension of the values array.

property components: List[Labels]

Get the component Labels for this block.

The entries in these labels describe intermediate dimensions of the values array.

property properties: Labels

Get the property Labels for this block.

The entries in these labels describe the last dimension of the values array. The properties are guaranteed to be the same for values and gradients in the same block.

copy() → TensorBlock

get a deep copy of this block, including all the data and metadata

Return type:

TensorBlock

add_gradient(parameter: str, gradient: TensorBlock)

Add gradient with respect to parameter in this block.

Parameters:

• parameter (str) – add gradients with respect to this parameter (e.g. positions, cell, …)

• gradient (TensorBlock) –

  a TensorBlock whose values contain the gradients of this TensorBlock values with respect to parameter. The labels of the gradient TensorBlock should be organized as follows:

  • its samples must contain "sample" as the first dimension, with values containing the index of the corresponding samples in this TensorBlock, and arbitrary supplementary samples dimension;

  • its components must contain at least the same components as this TensorBlock, with any additional components coming before those;

  • its properties must match exactly those of this TensorBlock.

>>> import numpy as np
>>> from metatensor.torch import TensorBlock, Labels
>>> block = TensorBlock(
...     values=torch.full((3, 1, 1), 1.0),
...     samples=Labels(["structure"], torch.IntTensor([[0], [2], [4]])),
...     components=[Labels.range("component", 1)],
...     properties=Labels.range("property", 1),
... )
>>> gradient = TensorBlock(
...     values=torch.full((2, 1, 1), 11.0),
...     samples=Labels(
...         names=["sample", "parameter"],
...         values=torch.IntTensor([[0, -2], [2, 3]]),
...     ),
...     components=[Labels.range("component", 1)],
...     properties=Labels.range("property", 1),
... )
>>> block.add_gradient("parameter", gradient)
>>> print(block)
TensorBlock
    samples (3): ['structure']
    components (1): ['component']
    properties (1): ['property']
    gradients: ['parameter']

gradient(parameter: str) → TensorBlock

Get the gradient of the block values with respect to the given parameter.

Parameters:

parameter (str) – check for gradients with respect to this parameter (e.g. positions, cell, …)

Return type:

TensorBlock

>>> from metatensor.torch import TensorBlock, Labels
>>> block = TensorBlock(
...     values=torch.full((3, 1, 5), 1.0),
...     samples=Labels(["structure"], torch.tensor([[0], [2], [4]])),
...     components=[Labels.range("component", 1)],
...     properties=Labels.range("property", 5),
... )

>>> positions_gradient = TensorBlock(
...     values=torch.full((2, 3, 1, 5), 11.0),
...     samples=Labels(["sample", "atom"], torch.tensor([[0, 2], [2, 3]])),
...     components=[
...         Labels.range("direction", 3),
...         Labels.range("component", 1),
...     ],
...     properties=Labels.range("property", 5),
... )
>>> block.add_gradient("positions", positions_gradient)

>>> cell_gradient = TensorBlock(
...     values=torch.full((2, 3, 3, 1, 5), 15.0),
...     samples=Labels.range("sample", 2),
...     components=[
...         Labels.range("direction_1", 3),
...         Labels.range("direction_2", 3),
...         Labels.range("component", 1),
...     ],
...     properties=Labels.range("property", 5),
... )
>>> block.add_gradient("cell", cell_gradient)

>>> positions_gradient = block.gradient("positions")
>>> print(positions_gradient)
Gradient TensorBlock ('positions')
    samples (2): ['sample', 'atom']
    components (3, 1): ['direction', 'component']
    properties (5): ['property']
    gradients: None

>>> cell_gradient = block.gradient("cell")
>>> print(cell_gradient)
Gradient TensorBlock ('cell')
    samples (2): ['sample']
    components (3, 3, 1): ['direction_1', 'direction_2', 'component']
    properties (5): ['property']
    gradients: None

gradients_list() → List[str]

get a list of all gradients defined in this block

Return type:

List[str]

has_gradient(parameter: str) → bool

Check if this block contains gradient information with respect to the given parameter.

Parameters:

parameter (str) – check for gradients with respect to this parameter (e.g. positions, cell, …)

Return type:

bool

gradients() → Dict[str, TensorBlock]

Get an iterator over all (parameter, gradients) pairs defined in this block.

Return type:

Dict[str, TensorBlock]