Source code for metatensor.operations.zeros_like

from typing import List, Optional, Union

from . import _dispatch
from ._backend import TensorBlock, TensorMap, torch_jit_script
from ._utils import _check_gradient_presence_raise


[docs] @torch_jit_script def zeros_like_block( block: TensorBlock, gradients: Optional[Union[List[str], str]] = None, requires_grad: bool = False, ) -> TensorBlock: """Return a new :py:class:`TensorBlock` with the same metadata as block, and all values equal to zero. :param block: Input block from which the metadata is taken. :param gradients: Which gradients should be present in the output. If this is :py:obj:`None` (default) all gradient of ``block`` are present in the new :py:class:`TensorBlock`. If this is an empty list ``[]``, no gradients information is copied. :param requires_grad: If autograd should record operations for the returned tensor. This option is only relevant for torch. """ values = _dispatch.zeros_like(block.values, requires_grad=requires_grad) result_block = TensorBlock( values=values, samples=block.samples, components=block.components, properties=block.properties, ) if isinstance(gradients, str): gradients = [gradients] if gradients is None: gradients = block.gradients_list() else: _check_gradient_presence_raise( block=block, parameters=gradients, fname="zeros_like" ) for parameter in gradients: gradient = block.gradient(parameter) if len(gradient.gradients_list()) != 0: raise NotImplementedError("gradients of gradients are not supported") gradient_values = _dispatch.zeros_like(gradient.values) result_block.add_gradient( parameter=parameter, gradient=TensorBlock( values=gradient_values, samples=gradient.samples, components=gradient.components, properties=gradient.properties, ), ) return result_block
[docs] @torch_jit_script def zeros_like( tensor: TensorMap, gradients: Optional[Union[List[str], str]] = None, requires_grad: bool = False, ) -> TensorMap: """Return a new :py:class:`TensorMap` with the same metadata as tensor, and all values equal to zero. :param tensor: Input tensor from which the metadata is taken. :param gradients: Which gradients should be present in the output. If this is :py:obj:`None` (default) all gradient of ``tensor`` are present in the new :py:class:`TensorMap`. If this is an empty list ``[]``, no gradients information is copied. :param requires_grad: If autograd should record operations for the returned tensor. This option is only relevant for torch. >>> import numpy as np >>> import metatensor >>> from metatensor import TensorBlock, TensorMap, Labels >>> np.random.seed(1) First we create a :py:class:`TensorMap` with just one block with two gradients, named ``alpha`` and ``beta``, containing random data: >>> block = TensorBlock( ... values=np.random.rand(4, 3), ... samples=Labels.range("sample", 4), ... components=[], ... properties=Labels.range("property", 3), ... ) >>> block.add_gradient( ... parameter="alpha", ... gradient=TensorBlock( ... values=np.random.rand(2, 3, 3), ... samples=Labels(["sample", "atom"], np.array([[0, 0], [0, 2]])), ... components=[Labels.range("component", 3)], ... properties=block.properties, ... ), ... ) >>> block.add_gradient( ... parameter="beta", ... gradient=TensorBlock( ... values=np.random.rand(1, 3), ... samples=Labels(["sample"], np.array([[0]])), ... components=[], ... properties=block.properties, ... ), ... ) >>> keys = Labels(names=["key"], values=np.array([[0]])) >>> tensor = TensorMap(keys, [block]) >>> print(tensor.block(0)) TensorBlock samples (4): ['sample'] components (): [] properties (3): ['property'] gradients: ['alpha', 'beta'] Then we use ``zeros_like`` to create a :py:class:`TensorMap` with the same metadata as ``tensor``, but with all values set to 0. >>> tensor_zeros = metatensor.zeros_like(tensor) >>> print(tensor_zeros.block(0)) TensorBlock samples (4): ['sample'] components (): [] properties (3): ['property'] gradients: ['alpha', 'beta'] >>> print(tensor_zeros.block(0).values) [[0. 0. 0.] [0. 0. 0.] [0. 0. 0.] [0. 0. 0.]] >>> print(tensor_zeros.block(0).gradient("alpha").values) [[[0. 0. 0.] [0. 0. 0.] [0. 0. 0.]] <BLANKLINE> [[0. 0. 0.] [0. 0. 0.] [0. 0. 0.]]] Note that if we copy just the gradient ``alpha``, ``beta`` is no longer available. >>> tensor_zeros = metatensor.zeros_like(tensor, gradients="alpha") >>> print(tensor_zeros.block(0).gradients_list()) ['alpha'] """ blocks: List[TensorBlock] = [] # this must be declared as a list of TensorBlocks for block in tensor.blocks(): blocks.append( zeros_like_block( block=block, gradients=gradients, requires_grad=requires_grad ) ) return TensorMap(tensor.keys, blocks)