TensorMap

using metatensor_torch::TorchTensorMap = torch::intrusive_ptr<TensorMapHolder>

TorchScript will always manipulate TensorMapHolder through a torch::intrusive_ptr

class TensorMapHolder : public CustomClassHolder

Wrapper around metatensor::TensorMap for integration with TorchScript

Python/TorchScript code will typically manipulate torch::intrusive_ptr<TensorMapHolder> (i.e. TorchTensorMap) instead of instances of TensorMapHolder.

Public Functions

TensorMapHolder(TorchLabels keys, const std::vector<TorchTensorBlock> &blocks)

Create a new TensorMapHolder for TorchScript.

In contrast to the TensorMap constructor, this does not move from the different blocks, but instead create new ones using the same data and metadata, but with incremented reference count.

TorchTensorMap copy() const

Make a copy of this TensorMap, including all the data contained inside.

TorchLabels keys() const

Get the keys for this TensorMap

std::vector<int64_t> blocks_matching(const TorchLabels &selection) const

Get a (possibly empty) list of block indexes matching the selection

TorchTensorMap keys_to_properties(torch::IValue keys_to_move, bool sort_samples) const

Merge blocks with the same value for selected keys dimensions along the property axis.

See metatensor::TensorMap::keys_to_properties for more information on this function.

The input torch::IValue can be a single string, a list/tuple of strings, or a TorchLabels instance.

TorchTensorMap keys_to_samples(torch::IValue keys_to_move, bool sort_samples) const

Merge blocks with the same value for selected keys dimensions along the sample axis.

See metatensor::TensorMap::keys_to_samples for more information on this function.

The input torch::IValue can be a single string, a list/tuple of strings, or a TorchLabels instance.

TorchTensorMap components_to_properties(torch::IValue dimensions) const

Move the given dimensions from the component labels to the property labels for each block.

See metatensor::TensorMap::components_to_properties for more information on this function.

The input torch::IValue can be a single string, or a list/tuple of strings.

std::vector<std::string> sample_names()

Get the names of the samples dimensions for all blocks in this TensorMap

std::vector<std::string> component_names()

Get the names of the components dimensions for all blocks in this TensorMap

std::vector<std::string> property_names()

Get the names of the properties dimensions for all blocks in this TensorMap

std::string print(int64_t max_keys) const

Print this TensorMap to a string, including at most max_keys in the output (-1 to include all keys).

torch::Device device() const

Get the device for the values stored in this TensorMap

torch::Dtype scalar_type() const

Get the dtype for the values stored in this TensorMap

TorchTensorMap to(torch::optional<torch::Dtype> dtype = torch::nullopt, torch::optional<torch::Device> device = torch::nullopt) const

Move this TensorMap to the given dtype and device.

TorchTensorMap to_positional(torch::IValue positional_1, torch::IValue positional_2, torch::optional<torch::Dtype> dtype, torch::optional<torch::Device> device, torch::optional<std::string> arrays) const

Wrapper of the to function to enable using it with positional parameters from Python; for example to(dtype), to(device), to(dtype, device=device), to(dtype, device), to(device, dtype), etc.

arrays is left as a keyword argument since it is mainly here for compatibility with the pure Python backend, and only "torch" is supported.

inline const metatensor::TensorMap &as_metatensor() const

Get the underlying metatensor TensorMap.

void save(const std::string &path) const

Serialize and save a TensorMap to the given path.

torch::Tensor save_buffer() const

Serialize and save a TensorMap to an in-memory buffer (represented as a torch::Tensor of bytes)

Public Static Functions

static TorchTensorBlock block_by_id(TorchTensorMap self, int64_t index)

Get a block inside this TensorMap by it’s index/the index of the corresponding key.

The returned TensorBlock is a view inside memory owned by this TensorMap, and is only valid as long as the TensorMap is kept alive.

static TorchTensorBlock block(TorchTensorMap self, const std::map<std::string, int32_t> &selection)

Get the block in this TensorMap with the key matching the name=>values passed in selection

static TorchTensorBlock block(TorchTensorMap self, TorchLabels selection)

Get the block in this TensorMap with the key matching the name=>values passed in selection. The selection must contain a single entry.

static TorchTensorBlock block(TorchTensorMap self, TorchLabelsEntry selection)

Get the block in this TensorMap with the key matching the name=>values passed in selection

static TorchTensorBlock block_torch(TorchTensorMap self, torch::IValue index)

TorchScript implementation of block, dispatching to one of the functions above

static std::vector<TorchTensorBlock> blocks_by_id(TorchTensorMap self, const std::vector<int64_t> &indices)

Similar to block_by_id, but get all blocks with the given indices.

static std::vector<TorchTensorBlock> blocks(TorchTensorMap self)

Get all blocks in this TensorMap.

static std::vector<TorchTensorBlock> blocks(TorchTensorMap self, const std::map<std::string, int32_t> &selection)

Similar to block, but allow getting multiple matching blocks.

static std::vector<TorchTensorBlock> blocks(TorchTensorMap self, TorchLabels selection)

Similar to block, but allow getting multiple matching blocks.

static std::vector<TorchTensorBlock> blocks(TorchTensorMap self, TorchLabelsEntry selection)

Similar to block, but allow getting multiple matching blocks.

static std::vector<TorchTensorBlock> blocks_torch(TorchTensorMap self, torch::IValue index)

TorchScript implementation of blocks, dispatching to one of the functions above.

static std::vector<std::tuple<TorchLabelsEntry, TorchTensorBlock>> items(TorchTensorMap self)

Get all (key => block) pairs in this TensorMap

static TorchTensorMap load(const std::string &path)

Load a serialized TensorMap from the given path.

static TorchTensorMap load_buffer(torch::Tensor buffer)

Load a serialized TensorMap from an in-memory buffer (represented as a torch::Tensor of bytes)