TensorMap

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

TorchScript will always manipulate TensorMapHolder through a torch::intrusive_ptr

class TensorMapHolder : public torch::CustomClassHolder

Wrapper around metatensor::TensorMap for integration with TorchScript

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

Public Functions

TensorMapHolder(Labels keys, std::vector<TensorBlock> blocks)

Create a new TensorMapHolder for TorchScript.

TensorMap copy() const

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

Labels keys() const

Get the keys for this TensorMap

TensorMap keys_to_properties(torch::IValue keys_to_move, torch::Scalar fill_value = 0, bool sort_samples = true) 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 Labels instance.

TensorMap keys_to_samples(torch::IValue keys_to_move, torch::Scalar fill_value = 0, bool sort_samples = true) 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 Labels instance.

TensorMap 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

TensorMap to(torch::optional<torch::Dtype> dtype = torch::nullopt, torch::optional<torch::Device> device = torch::nullopt, bool non_blocking = false) const

Move this TensorMap to the given dtype and device.

TensorMap 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, bool non_blocking = false) 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)

void set_info(std::string key, std::string value)

Set or update the info (i.e. global metadata) value associated with key for this TensorMap.

torch::optional<std::string> get_info(std::string key) const

Get the info (i.e. global metadata) with the given key for this TensorMap.

torch::Dict<std::string, std::string> info() const

Get all the key/value info pairs stored in this TensorMap.

Public Static Functions

static TensorBlock block_by_id(TensorMap 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 TensorBlock block(TensorMap 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 TensorBlock block(TensorMap self, Labels 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 TensorBlock block(TensorMap self, LabelsEntry selection)

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

static TensorBlock block_torch(TensorMap self, torch::IValue index)

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

static std::vector<TensorBlock> blocks_by_id(TensorMap self, const std::vector<int64_t> &indices)

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

static std::vector<TensorBlock> blocks(TensorMap self)

Get all blocks in this TensorMap.

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

Similar to block, but allow getting multiple matching blocks.

static std::vector<TensorBlock> blocks(TensorMap self, Labels selection)

Similar to block, but allow getting multiple matching blocks.

static std::vector<TensorBlock> blocks(TensorMap self, LabelsEntry selection)

Similar to block, but allow getting multiple matching blocks.

static std::vector<TensorBlock> blocks_torch(TensorMap self, torch::IValue index)

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

static std::vector<std::tuple<LabelsEntry, TensorBlock>> items(TensorMap self)

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

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

Load a serialized TensorMap from the given path.

static TensorMap load_buffer(torch::Tensor buffer)

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