From d7ba0fd3915ae4b08b0c9e9a1d0339b666b0a926 Mon Sep 17 00:00:00 2001 From: SanggyuChong Date: Sat, 17 Feb 2024 15:30:56 +0100 Subject: [PATCH 1/7] implement contiguous functions within temporary ipynb --- .../metatensor/operations/_dispatch.py | 36 +++ .../operations/temp_implementations.ipynb | 224 ++++++++++++++++++ 2 files changed, 260 insertions(+) create mode 100644 python/metatensor-operations/metatensor/operations/temp_implementations.ipynb diff --git a/python/metatensor-operations/metatensor/operations/_dispatch.py b/python/metatensor-operations/metatensor/operations/_dispatch.py index 7dd2773fa..88ede1920 100644 --- a/python/metatensor-operations/metatensor/operations/_dispatch.py +++ b/python/metatensor-operations/metatensor/operations/_dispatch.py @@ -343,6 +343,42 @@ def concatenate(arrays: List[TorchTensor], axis: int): raise TypeError(UNKNOWN_ARRAY_TYPE) +def is_contiguous_array(array): + """ + Checks if a given array is contiguous. + + In the case of numpy, C order is used for consistency with torch. As such, + only C-contiguity is checked. + """ + if isinstance(array, TorchTensor): + return array.is_contiguous() + + elif isinstance(array, np.ndarray): + return array.flags["C_CONTIGUOUS"] + + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + +def make_contiguous_array(array): + """ + Returns a contiguous array. + It is equivalent of np.ascontiguousarray(array) and tensor.contiguous(). In + the case of numpy, C order is used for consistency with torch. As such, only + C-contiguity is checked. + """ + if isinstance(array, TorchTensor): + if array.is_contiguous(): + return array + return array.contiguous() + elif isinstance(array, np.ndarray): + if array.flags["C_CONTIGUOUS"]: + return array + return np.ascontiguousarray(array) + else: + raise TypeError(UNKNOWN_ARRAY_TYPE) + + def copy(array): """Returns a copy of ``array``. The new data is not shared with the original array""" diff --git a/python/metatensor-operations/metatensor/operations/temp_implementations.ipynb b/python/metatensor-operations/metatensor/operations/temp_implementations.ipynb new file mode 100644 index 000000000..45baf73f6 --- /dev/null +++ b/python/metatensor-operations/metatensor/operations/temp_implementations.ipynb @@ -0,0 +1,224 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 24, + "metadata": {}, + "outputs": [], + "source": [ + "import metatensor as mts\n", + "import numpy as np\n", + "\n", + "from _dispatch import is_contiguous_array, make_contiguous_array\n", + "from metatensor import TensorMap, TensorBlock, Labels\n", + "\n", + "# 3. Do the same for make_contiguous{_block}\n", + "\n", + "tensor = mts.load(\"/Users/sanggyu/Research/metatensor/metatensor/python/metatensor-operations/tests/data/qm7-power-spectrum.npz\", use_numpy=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": {}, + "outputs": [], + "source": [ + "def make_incontiguous_block(block: TensorBlock) -> TensorBlock:\n", + " \n", + " \"\"\"\n", + " Make a non-contiguous block by reversing the order in both the main value\n", + " block and the gradient block(s). \n", + " \n", + " TODO -- make it applicable to all blocks, not just blocks with more than one\n", + " row\n", + " \"\"\"\n", + " \n", + " new_block = TensorBlock(\n", + " values=block.values.copy()[::-1],\n", + " samples=block.samples,\n", + " components=block.components,\n", + " properties=block.properties,\n", + " )\n", + " # Make gradients non-contig\n", + " for param, gradient in block.gradients():\n", + " \n", + " new_gradient = TensorBlock(\n", + " values=gradient.values.copy()[::-1],\n", + " samples=gradient.samples,\n", + " components=gradient.components,\n", + " properties=gradient.properties,\n", + " )\n", + " new_block.add_gradient(param, new_gradient)\n", + " \n", + " return new_block" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "def make_incontiguous(tensor: TensorMap) -> TensorMap:\n", + " \n", + " \"\"\"\n", + " Make a non-contiguous TensorMap by reversing the order in all the main value\n", + " blocks and the gradient blocks. \n", + " \"\"\"\n", + " \n", + " keys = tensor.keys\n", + " new_blocks = []\n", + " \n", + " for key, block in tensor.items():\n", + " # Create a new TM with a non-contig array\n", + " new_block = make_incontiguous_block(block)\n", + " new_blocks.append(new_block)\n", + " \n", + " return TensorMap(keys=keys, blocks=new_blocks)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "def is_contiguous_block(block: TensorBlock) -> bool:\n", + " \n", + " \"\"\"PUBLIC FUNCTION TO PUBLISH\"\"\"\n", + " \n", + " check_contiguous = True\n", + " if not is_contiguous_array(block.values):\n", + " check_contiguous = False\n", + " \n", + " return check_contiguous" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "def is_contiguous(tensor: TensorMap) -> bool:\n", + " \n", + " \"\"\"PUBLIC FUNCTION TO PUBLISH\"\"\"\n", + " \n", + " check_contiguous = True\n", + " for key, block in tensor.items():\n", + " # Here, call another function: def is_contiguous_block(block: TensorBlock) -> bool\n", + " if not is_contiguous_block(block):\n", + " check_contiguous = False\n", + " \n", + " for param, gradient in block.gradients():\n", + " if not is_contiguous_block(gradient):\n", + " check_contiguous = False\n", + " \n", + " return check_contiguous" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "metadata": {}, + "outputs": [], + "source": [ + "def make_contiguous_block(block: TensorBlock) -> TensorBlock:\n", + " \n", + " \"\"\"PUBLIC FUNCTION TO PUBLISH\"\"\"\n", + " \n", + " contiguous_block = TensorBlock(\n", + " values=make_contiguous_array(block.values.copy()),\n", + " samples=block.samples,\n", + " components=block.components,\n", + " properties=block.properties,\n", + " )\n", + " # Make gradients non-contig\n", + " for param, gradient in block.gradients():\n", + "\n", + " new_gradient = TensorBlock(\n", + " values=make_contiguous_array(gradient.values.copy()),\n", + " samples=gradient.samples,\n", + " components=gradient.components,\n", + " properties=gradient.properties,\n", + " )\n", + " contiguous_block.add_gradient(param, new_gradient)\n", + " \n", + " return contiguous_block" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "metadata": {}, + "outputs": [], + "source": [ + "def make_contiguous(tensor: TensorMap) -> TensorMap:\n", + " \n", + " \"\"\"PUBLIC FUNCTION TO PUBLISH\"\"\"\n", + " \n", + " keys = tensor.keys\n", + " contiguous_blocks = []\n", + " \n", + " for key, block in tensor.items():\n", + " # Create a new TM with a non-contig array\n", + " contiguous_block = make_contiguous_block(block)\n", + " contiguous_blocks.append(contiguous_block)\n", + " \n", + " return TensorMap(keys=keys, blocks=contiguous_blocks)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "metadata": {}, + "outputs": [], + "source": [ + "incontig_tensor = make_incontiguous(tensor)\n", + "assert is_contiguous(tensor)\n", + "assert not is_contiguous(incontig_tensor)\n", + "assert is_contiguous(make_contiguous(incontig_tensor))" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "metadata": {}, + "outputs": [], + "source": [ + "incontig_block = make_incontiguous_block(tensor.block(0))\n", + "assert is_contiguous_block(tensor.block(0))\n", + "assert not is_contiguous_block(incontig_block)\n", + "assert is_contiguous_block(make_contiguous_block(incontig_block))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "mts-dev", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} From 7b737ab42496cfc72da2f88c007768ede9fb2caf Mon Sep 17 00:00:00 2001 From: SanggyuChong Date: Tue, 20 Feb 2024 16:31:37 +0100 Subject: [PATCH 2/7] implement is_contiguous and make_contiguous into actual operations, and write tests for them --- .../metatensor/operations/__init__.py | 6 + .../metatensor/operations/contiguous.py | 91 +++++++ .../operations/temp_implementations.ipynb | 224 ------------------ .../metatensor-operations/tests/contiguous.py | 71 ++++++ 4 files changed, 168 insertions(+), 224 deletions(-) create mode 100644 python/metatensor-operations/metatensor/operations/contiguous.py delete mode 100644 python/metatensor-operations/metatensor/operations/temp_implementations.ipynb create mode 100644 python/metatensor-operations/tests/contiguous.py diff --git a/python/metatensor-operations/metatensor/operations/__init__.py b/python/metatensor-operations/metatensor/operations/__init__.py index d37d5fb87..7c8d413b9 100644 --- a/python/metatensor-operations/metatensor/operations/__init__.py +++ b/python/metatensor-operations/metatensor/operations/__init__.py @@ -15,6 +15,12 @@ allclose_raise, ) from .block_from_array import block_from_array # noqa +from .contiguous import ( + is_contiguous, + is_contiguous_block, + make_contiguous, + make_contiguous_block +) from .detach import detach, detach_block # noqa from .divide import divide # noqa from .dot import dot # noqa diff --git a/python/metatensor-operations/metatensor/operations/contiguous.py b/python/metatensor-operations/metatensor/operations/contiguous.py new file mode 100644 index 000000000..c42fda490 --- /dev/null +++ b/python/metatensor-operations/metatensor/operations/contiguous.py @@ -0,0 +1,91 @@ +from . import _dispatch +from ._backend import ( + Labels, + TensorBlock, + TensorMap, + torch_jit_is_scripting, + torch_jit_script, +) + +def is_contiguous(tensor: TensorMap) -> bool: + r""" + Checks contiguity of values and gradients in a :py:class:`TensorMap` + + :param A: the input :py:class:`TensorMap`. + + :return: a boolean indicating (non-)contiguity of values in ``A``. + """ + check_contiguous = True + for key, block in tensor.items(): + # Here, call another function: def is_contiguous_block(block: TensorBlock) -> bool + if not is_contiguous_block(block): + check_contiguous = False + + for param, gradient in block.gradients(): + if not is_contiguous_block(gradient): + check_contiguous = False + + return check_contiguous + + +def is_contiguous_block(block: TensorBlock) -> bool: + r""" + Checks contiguity of values in a :py:class:`TensorBlock` + + :param A: the input :py:class:`TensorBlock`. + + :return: a boolean indicating (non-)contiguity of values in ``A``. + """ + check_contiguous = True + if not _dispatch.is_contiguous_array(block.values): + check_contiguous = False + + return check_contiguous + + +def make_contiguous(tensor: TensorMap) -> TensorMap: + r""" + Return a new :py:class:`TensorMap` where the values are made to be contiguous. + + :param A: the input :py:class:`TensorMap`. + + :return: a new :py:class:`TensorMap` with the same metadata as ``A`` and + contiguous values of ``A``. + """ + + keys = tensor.keys + contiguous_blocks = [] + + for key, block in tensor.items(): + contiguous_block = make_contiguous_block(block) + contiguous_blocks.append(contiguous_block) + + return TensorMap(keys=keys, blocks=contiguous_blocks) + + +def make_contiguous_block(block: TensorBlock) -> TensorBlock: + r""" + Return a new :py:class:`TensorBlock` where the values are made to be contiguous. + + :param A: the input :py:class:`TensorBlock`. + + :return: a new :py:class:`TensorBlock` where values are contiguous. + """ + contiguous_block = TensorBlock( + values=_dispatch.make_contiguous_array(block.values.copy()), + samples=block.samples, + components=block.components, + properties=block.properties, + ) + # Make gradients non-contig + for param, gradient in block.gradients(): + + new_gradient = TensorBlock( + values=_dispatch.make_contiguous_array(gradient.values.copy()), + samples=gradient.samples, + components=gradient.components, + properties=gradient.properties, + ) + contiguous_block.add_gradient(param, new_gradient) + + return contiguous_block \ No newline at end of file diff --git a/python/metatensor-operations/metatensor/operations/temp_implementations.ipynb b/python/metatensor-operations/metatensor/operations/temp_implementations.ipynb deleted file mode 100644 index 45baf73f6..000000000 --- a/python/metatensor-operations/metatensor/operations/temp_implementations.ipynb +++ /dev/null @@ -1,224 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 24, - "metadata": {}, - "outputs": [], - "source": [ - "import metatensor as mts\n", - "import numpy as np\n", - "\n", - "from _dispatch import is_contiguous_array, make_contiguous_array\n", - "from metatensor import TensorMap, TensorBlock, Labels\n", - "\n", - "# 3. Do the same for make_contiguous{_block}\n", - "\n", - "tensor = mts.load(\"/Users/sanggyu/Research/metatensor/metatensor/python/metatensor-operations/tests/data/qm7-power-spectrum.npz\", use_numpy=True)" - ] - }, - { - "cell_type": "code", - "execution_count": 21, - "metadata": {}, - "outputs": [], - "source": [ - "def make_incontiguous_block(block: TensorBlock) -> TensorBlock:\n", - " \n", - " \"\"\"\n", - " Make a non-contiguous block by reversing the order in both the main value\n", - " block and the gradient block(s). \n", - " \n", - " TODO -- make it applicable to all blocks, not just blocks with more than one\n", - " row\n", - " \"\"\"\n", - " \n", - " new_block = TensorBlock(\n", - " values=block.values.copy()[::-1],\n", - " samples=block.samples,\n", - " components=block.components,\n", - " properties=block.properties,\n", - " )\n", - " # Make gradients non-contig\n", - " for param, gradient in block.gradients():\n", - " \n", - " new_gradient = TensorBlock(\n", - " values=gradient.values.copy()[::-1],\n", - " samples=gradient.samples,\n", - " components=gradient.components,\n", - " properties=gradient.properties,\n", - " )\n", - " new_block.add_gradient(param, new_gradient)\n", - " \n", - " return new_block" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [], - "source": [ - "def make_incontiguous(tensor: TensorMap) -> TensorMap:\n", - " \n", - " \"\"\"\n", - " Make a non-contiguous TensorMap by reversing the order in all the main value\n", - " blocks and the gradient blocks. \n", - " \"\"\"\n", - " \n", - " keys = tensor.keys\n", - " new_blocks = []\n", - " \n", - " for key, block in tensor.items():\n", - " # Create a new TM with a non-contig array\n", - " new_block = make_incontiguous_block(block)\n", - " new_blocks.append(new_block)\n", - " \n", - " return TensorMap(keys=keys, blocks=new_blocks)" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [], - "source": [ - "def is_contiguous_block(block: TensorBlock) -> bool:\n", - " \n", - " \"\"\"PUBLIC FUNCTION TO PUBLISH\"\"\"\n", - " \n", - " check_contiguous = True\n", - " if not is_contiguous_array(block.values):\n", - " check_contiguous = False\n", - " \n", - " return check_contiguous" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": {}, - "outputs": [], - "source": [ - "def is_contiguous(tensor: TensorMap) -> bool:\n", - " \n", - " \"\"\"PUBLIC FUNCTION TO PUBLISH\"\"\"\n", - " \n", - " check_contiguous = True\n", - " for key, block in tensor.items():\n", - " # Here, call another function: def is_contiguous_block(block: TensorBlock) -> bool\n", - " if not is_contiguous_block(block):\n", - " check_contiguous = False\n", - " \n", - " for param, gradient in block.gradients():\n", - " if not is_contiguous_block(gradient):\n", - " check_contiguous = False\n", - " \n", - " return check_contiguous" - ] - }, - { - "cell_type": "code", - "execution_count": 25, - "metadata": {}, - "outputs": [], - "source": [ - "def make_contiguous_block(block: TensorBlock) -> TensorBlock:\n", - " \n", - " \"\"\"PUBLIC FUNCTION TO PUBLISH\"\"\"\n", - " \n", - " contiguous_block = TensorBlock(\n", - " values=make_contiguous_array(block.values.copy()),\n", - " samples=block.samples,\n", - " components=block.components,\n", - " properties=block.properties,\n", - " )\n", - " # Make gradients non-contig\n", - " for param, gradient in block.gradients():\n", - "\n", - " new_gradient = TensorBlock(\n", - " values=make_contiguous_array(gradient.values.copy()),\n", - " samples=gradient.samples,\n", - " components=gradient.components,\n", - " properties=gradient.properties,\n", - " )\n", - " contiguous_block.add_gradient(param, new_gradient)\n", - " \n", - " return contiguous_block" - ] - }, - { - "cell_type": "code", - "execution_count": 26, - "metadata": {}, - "outputs": [], - "source": [ - "def make_contiguous(tensor: TensorMap) -> TensorMap:\n", - " \n", - " \"\"\"PUBLIC FUNCTION TO PUBLISH\"\"\"\n", - " \n", - " keys = tensor.keys\n", - " contiguous_blocks = []\n", - " \n", - " for key, block in tensor.items():\n", - " # Create a new TM with a non-contig array\n", - " contiguous_block = make_contiguous_block(block)\n", - " contiguous_blocks.append(contiguous_block)\n", - " \n", - " return TensorMap(keys=keys, blocks=contiguous_blocks)\n" - ] - }, - { - "cell_type": "code", - "execution_count": 29, - "metadata": {}, - "outputs": [], - "source": [ - "incontig_tensor = make_incontiguous(tensor)\n", - "assert is_contiguous(tensor)\n", - "assert not is_contiguous(incontig_tensor)\n", - "assert is_contiguous(make_contiguous(incontig_tensor))" - ] - }, - { - "cell_type": "code", - "execution_count": 30, - "metadata": {}, - "outputs": [], - "source": [ - "incontig_block = make_incontiguous_block(tensor.block(0))\n", - "assert is_contiguous_block(tensor.block(0))\n", - "assert not is_contiguous_block(incontig_block)\n", - "assert is_contiguous_block(make_contiguous_block(incontig_block))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "mts-dev", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.0" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/python/metatensor-operations/tests/contiguous.py b/python/metatensor-operations/tests/contiguous.py new file mode 100644 index 000000000..89efa1f25 --- /dev/null +++ b/python/metatensor-operations/tests/contiguous.py @@ -0,0 +1,71 @@ +import os + +import metatensor +from metatensor import TensorBlock, TensorMap + +DATA_ROOT = os.path.join(os.path.dirname(__file__), "data") + +tensor = metatensor.load( + os.path.join(DATA_ROOT, "qm7-power-spectrum.npz"), + use_numpy=True, + ) + + +def _incontiguous_tensor(tensor: TensorMap) -> TensorMap: + """ + Make a non-contiguous TensorMap by reversing the order in all the main value + blocks and the gradient blocks. + """ + + keys = tensor.keys + new_blocks = [] + + for key, block in tensor.items(): + # Create a new TM with a non-contig array + new_block = _incontiguous_block(block) + new_blocks.append(new_block) + + return TensorMap(keys=keys, blocks=new_blocks) + + +def _incontiguous_block(block: TensorBlock) -> TensorBlock: + """ + Make a non-contiguous block by reversing the order in both the main value block and + the gradient block(s). + """ + new_vals = block.values.copy()[::-1, ::-1] + new_block = TensorBlock( + values=new_vals, + samples=block.samples, + components=block.components, + properties=block.properties, + ) + # Make gradients non-contig + for param, gradient in block.gradients(): + new_grad = gradient.values.copy()[::-1, ::-1] + new_gradient = TensorBlock( + values=new_grad, + samples=gradient.samples, + components=gradient.components, + properties=gradient.properties, + ) + new_block.add_gradient(param, new_gradient) + + return new_block + + +def test_is_contiguous_block(): + assert not metatensor.is_contiguous_block(_incontiguous_block(tensor.block(0))) + assert metatensor.is_contiguous_block( + metatensor.make_contiguous_block(_incontiguous_block(tensor.block(0))) + ) + + +def test_is_contiguous(): + assert not metatensor.is_contiguous(_incontiguous_tensor(tensor)) + assert metatensor.is_contiguous( + metatensor.make_contiguous(_incontiguous_tensor(tensor)) + ) + + + \ No newline at end of file From f8929d956b3dcf34aa0a2ab3e99d1b3de7451b49 Mon Sep 17 00:00:00 2001 From: SanggyuChong Date: Mon, 26 Feb 2024 18:00:41 +0100 Subject: [PATCH 3/7] linting, formatting, responding to review comments --- .../metatensor/operations/__init__.py | 10 ++-- .../metatensor/operations/contiguous.py | 52 +++++++++---------- .../metatensor-operations/tests/contiguous.py | 35 +++++++------ 3 files changed, 49 insertions(+), 48 deletions(-) diff --git a/python/metatensor-operations/metatensor/operations/__init__.py b/python/metatensor-operations/metatensor/operations/__init__.py index 7c8d413b9..c7adcc509 100644 --- a/python/metatensor-operations/metatensor/operations/__init__.py +++ b/python/metatensor-operations/metatensor/operations/__init__.py @@ -15,11 +15,11 @@ allclose_raise, ) from .block_from_array import block_from_array # noqa -from .contiguous import ( - is_contiguous, - is_contiguous_block, - make_contiguous, - make_contiguous_block +from .contiguous import ( # noqa: F401 + is_contiguous, + is_contiguous_block, + make_contiguous, + make_contiguous_block, ) from .detach import detach, detach_block # noqa from .divide import divide # noqa diff --git a/python/metatensor-operations/metatensor/operations/contiguous.py b/python/metatensor-operations/metatensor/operations/contiguous.py index c42fda490..53538699b 100644 --- a/python/metatensor-operations/metatensor/operations/contiguous.py +++ b/python/metatensor-operations/metatensor/operations/contiguous.py @@ -1,27 +1,26 @@ from . import _dispatch -from ._backend import ( - Labels, +from ._backend import ( # torch_jit_is_scripting,; torch_jit_script, TensorBlock, TensorMap, - torch_jit_is_scripting, - torch_jit_script, ) + def is_contiguous(tensor: TensorMap) -> bool: - r""" + """ Checks contiguity of values and gradients in a :py:class:`TensorMap` - - :param A: the input :py:class:`TensorMap`. - :return: a boolean indicating (non-)contiguity of values in ``A``. + :param tensor: the input :py:class:`TensorMap`. + + :return: a boolean indicating (non-)contiguity of values in ``tensor``. Returns true + if all values and gradients arrays in the input ``tensor`` are contiguous, false + otherwise. """ check_contiguous = True - for key, block in tensor.items(): - # Here, call another function: def is_contiguous_block(block: TensorBlock) -> bool + for _key, block in tensor.items(): if not is_contiguous_block(block): check_contiguous = False - for param, gradient in block.gradients(): + for _param, gradient in block.gradients(): if not is_contiguous_block(gradient): check_contiguous = False @@ -29,12 +28,11 @@ def is_contiguous(tensor: TensorMap) -> bool: def is_contiguous_block(block: TensorBlock) -> bool: - r""" - Checks contiguity of values in a :py:class:`TensorBlock` - - :param A: the input :py:class:`TensorBlock`. + """ + :param block: the input :py:class:`TensorBlock`. - :return: a boolean indicating (non-)contiguity of values in ``A``. + :return: a boolean indicating (non-)contiguity of values in `TensorBlock` supplied + as input. """ check_contiguous = True if not _dispatch.is_contiguous_array(block.values): @@ -44,19 +42,19 @@ def is_contiguous_block(block: TensorBlock) -> bool: def make_contiguous(tensor: TensorMap) -> TensorMap: - r""" + """ Return a new :py:class:`TensorMap` where the values are made to be contiguous. - :param A: the input :py:class:`TensorMap`. + :param tensor: the input :py:class:`TensorMap`. - :return: a new :py:class:`TensorMap` with the same metadata as ``A`` and - contiguous values of ``A``. + :return: a new :py:class:`TensorMap` with the same data and metadata as ``tensor`` + and contiguous values of ``tensor``. """ - + keys = tensor.keys contiguous_blocks = [] - for key, block in tensor.items(): + for _key, block in tensor.items(): contiguous_block = make_contiguous_block(block) contiguous_blocks.append(contiguous_block) @@ -64,12 +62,13 @@ def make_contiguous(tensor: TensorMap) -> TensorMap: def make_contiguous_block(block: TensorBlock) -> TensorBlock: - r""" + """ Return a new :py:class:`TensorBlock` where the values are made to be contiguous. - :param A: the input :py:class:`TensorBlock`. + :param block: the input :py:class:`TensorBlock`. - :return: a new :py:class:`TensorBlock` where values are contiguous. + :return: a new :py:class:`TensorBlock` where all the values and gradients arrays (if + present) are contiguous. """ contiguous_block = TensorBlock( values=_dispatch.make_contiguous_array(block.values.copy()), @@ -77,6 +76,7 @@ def make_contiguous_block(block: TensorBlock) -> TensorBlock: components=block.components, properties=block.properties, ) + # Make gradients non-contig for param, gradient in block.gradients(): @@ -88,4 +88,4 @@ def make_contiguous_block(block: TensorBlock) -> TensorBlock: ) contiguous_block.add_gradient(param, new_gradient) - return contiguous_block \ No newline at end of file + return contiguous_block diff --git a/python/metatensor-operations/tests/contiguous.py b/python/metatensor-operations/tests/contiguous.py index 89efa1f25..4b680d4cc 100644 --- a/python/metatensor-operations/tests/contiguous.py +++ b/python/metatensor-operations/tests/contiguous.py @@ -1,26 +1,33 @@ import os +import pytest + import metatensor from metatensor import TensorBlock, TensorMap + DATA_ROOT = os.path.join(os.path.dirname(__file__), "data") -tensor = metatensor.load( + +@pytest.fixture +def tensor(): + """Loads a TensorMap from file for use in tests""" + return metatensor.load( os.path.join(DATA_ROOT, "qm7-power-spectrum.npz"), use_numpy=True, ) -def _incontiguous_tensor(tensor: TensorMap) -> TensorMap: +@pytest.fixture +def incontiguous_tensor(tensor) -> TensorMap: """ - Make a non-contiguous TensorMap by reversing the order in all the main value - blocks and the gradient blocks. + Make a TensorMap non-contiguous by reversing the order of the samples/properties + rows/columns in all values and gradient blocks. """ - keys = tensor.keys new_blocks = [] - for key, block in tensor.items(): + for _key, block in tensor.items(): # Create a new TM with a non-contig array new_block = _incontiguous_block(block) new_blocks.append(new_block) @@ -40,7 +47,6 @@ def _incontiguous_block(block: TensorBlock) -> TensorBlock: components=block.components, properties=block.properties, ) - # Make gradients non-contig for param, gradient in block.gradients(): new_grad = gradient.values.copy()[::-1, ::-1] new_gradient = TensorBlock( @@ -54,18 +60,13 @@ def _incontiguous_block(block: TensorBlock) -> TensorBlock: return new_block -def test_is_contiguous_block(): +def test_is_contiguous_block(tensor): assert not metatensor.is_contiguous_block(_incontiguous_block(tensor.block(0))) assert metatensor.is_contiguous_block( metatensor.make_contiguous_block(_incontiguous_block(tensor.block(0))) - ) + ) -def test_is_contiguous(): - assert not metatensor.is_contiguous(_incontiguous_tensor(tensor)) - assert metatensor.is_contiguous( - metatensor.make_contiguous(_incontiguous_tensor(tensor)) - ) - - - \ No newline at end of file +def test_is_contiguous(incontiguous_tensor): + assert not metatensor.is_contiguous(incontiguous_tensor) + assert metatensor.is_contiguous(metatensor.make_contiguous(incontiguous_tensor)) From 03501bf594aa98505939e93eaaea0a876bf4ed49 Mon Sep 17 00:00:00 2001 From: Joseph Abbott Date: Wed, 3 Apr 2024 17:44:56 +0200 Subject: [PATCH 4/7] Clean up, attempt torchscript version --- .../metatensor/operations/contiguous.py | 91 ++++++++++--------- .../metatensor-operations/tests/contiguous.py | 2 +- .../tests/operations/contiguous.py | 10 ++ 3 files changed, 61 insertions(+), 42 deletions(-) create mode 100644 python/metatensor-torch/tests/operations/contiguous.py diff --git a/python/metatensor-operations/metatensor/operations/contiguous.py b/python/metatensor-operations/metatensor/operations/contiguous.py index 53538699b..e9996d4e6 100644 --- a/python/metatensor-operations/metatensor/operations/contiguous.py +++ b/python/metatensor-operations/metatensor/operations/contiguous.py @@ -1,73 +1,65 @@ from . import _dispatch -from ._backend import ( # torch_jit_is_scripting,; torch_jit_script, +from ._backend import ( + Labels, TensorBlock, TensorMap, + torch_jit_is_scripting, + torch_jit_script, ) -def is_contiguous(tensor: TensorMap) -> bool: - """ - Checks contiguity of values and gradients in a :py:class:`TensorMap` - - :param tensor: the input :py:class:`TensorMap`. - - :return: a boolean indicating (non-)contiguity of values in ``tensor``. Returns true - if all values and gradients arrays in the input ``tensor`` are contiguous, false - otherwise. +@torch_jit_script +def is_contiguous_block(block: TensorBlock) -> bool: """ - check_contiguous = True - for _key, block in tensor.items(): - if not is_contiguous_block(block): - check_contiguous = False - - for _param, gradient in block.gradients(): - if not is_contiguous_block(gradient): - check_contiguous = False - - return check_contiguous + Checks whether the values array and gradients values arrays (if present) of an input + :py:class:`TensorBlock` are contiguous. + Note that arrays of :py:class:`Labels` objects are not checked for contiguity. -def is_contiguous_block(block: TensorBlock) -> bool: - """ :param block: the input :py:class:`TensorBlock`. - :return: a boolean indicating (non-)contiguity of values in `TensorBlock` supplied - as input. + :return: bool, true if all values arrays contiguous, false otherwise. """ - check_contiguous = True + check_contiguous: bool = True if not _dispatch.is_contiguous_array(block.values): check_contiguous = False + for _param, gradient in block.gradients(): + if not is_contiguous_block(gradient): + check_contiguous = False + return check_contiguous -def make_contiguous(tensor: TensorMap) -> TensorMap: +@torch_jit_script +def is_contiguous(tensor: TensorMap) -> bool: """ - Return a new :py:class:`TensorMap` where the values are made to be contiguous. + Checks whether all values arrays and gradients values arrays (if present) in all + :py:class:`TensorBlock` of an input :py:class:`TensorMap` are contiguous. + + Note that arrays of :py:class:`Labels` objects are not checked for contiguity. :param tensor: the input :py:class:`TensorMap`. - :return: a new :py:class:`TensorMap` with the same data and metadata as ``tensor`` - and contiguous values of ``tensor``. + :return: bool, true if all values arrays contiguous, false otherwise. """ - - keys = tensor.keys - contiguous_blocks = [] - + check_contiguous: bool = True for _key, block in tensor.items(): - contiguous_block = make_contiguous_block(block) - contiguous_blocks.append(contiguous_block) + if not is_contiguous_block(block): + check_contiguous = False - return TensorMap(keys=keys, blocks=contiguous_blocks) + return check_contiguous +@torch_jit_script def make_contiguous_block(block: TensorBlock) -> TensorBlock: """ - Return a new :py:class:`TensorBlock` where the values are made to be contiguous. + Returns a new :py:class:`TensorBlock` where the values and gradient values (if + present) arrays are mades to be contiguous. :param block: the input :py:class:`TensorBlock`. - :return: a new :py:class:`TensorBlock` where all the values and gradients arrays (if + :return: a new :py:class:`TensorBlock` where the values and gradients arrays (if present) are contiguous. """ contiguous_block = TensorBlock( @@ -76,10 +68,7 @@ def make_contiguous_block(block: TensorBlock) -> TensorBlock: components=block.components, properties=block.properties, ) - - # Make gradients non-contig for param, gradient in block.gradients(): - new_gradient = TensorBlock( values=_dispatch.make_contiguous_array(gradient.values.copy()), samples=gradient.samples, @@ -89,3 +78,23 @@ def make_contiguous_block(block: TensorBlock) -> TensorBlock: contiguous_block.add_gradient(param, new_gradient) return contiguous_block + + +@torch_jit_script +def make_contiguous(tensor: TensorMap) -> TensorMap: + """ + Returns a new :py:class:`TensorMap` where all values and gradient values arrays are + mades to be contiguous. + + :param tensor: the input :py:class:`TensorMap`. + + :return: a new :py:class:`TensorMap` with the same data and metadata as ``tensor`` + and contiguous values of ``tensor``. + """ + keys: Labels = tensor.keys + contiguous_blocks: List[TensorBlock] = [] + for _key, block in tensor.items(): + contiguous_block = make_contiguous_block(block) + contiguous_blocks.append(contiguous_block) + + return TensorMap(keys=keys, blocks=contiguous_blocks) diff --git a/python/metatensor-operations/tests/contiguous.py b/python/metatensor-operations/tests/contiguous.py index 4b680d4cc..2f66d8a13 100644 --- a/python/metatensor-operations/tests/contiguous.py +++ b/python/metatensor-operations/tests/contiguous.py @@ -69,4 +69,4 @@ def test_is_contiguous_block(tensor): def test_is_contiguous(incontiguous_tensor): assert not metatensor.is_contiguous(incontiguous_tensor) - assert metatensor.is_contiguous(metatensor.make_contiguous(incontiguous_tensor)) + assert metatensor.is_contiguous(metatensor.make_contiguous(incontiguous_tensor)) \ No newline at end of file diff --git a/python/metatensor-torch/tests/operations/contiguous.py b/python/metatensor-torch/tests/operations/contiguous.py new file mode 100644 index 000000000..c6919a390 --- /dev/null +++ b/python/metatensor-torch/tests/operations/contiguous.py @@ -0,0 +1,10 @@ +import io + +import torch + +import metatensor.torch + + +def test_is_contiguous(): + # TODO: write tests, used as a placeholder for now + assert True From de733bea9efc1c208f5dc95c932dd78c6571fec7 Mon Sep 17 00:00:00 2001 From: Joseph Abbott Date: Wed, 3 Apr 2024 18:02:10 +0200 Subject: [PATCH 5/7] Separate out is_contiguous (logic) and make_contiguous (manipulation) opeartions. Add docs --- docs/src/operations/reference/logic/index.rst | 1 + .../reference/logic/is_contiguous.rst | 7 ++ .../reference/manipulation/index.rst | 1 + .../manipulation/make_contiguous.rst | 6 ++ .../metatensor/operations/__init__.py | 14 ++-- .../metatensor/operations/is_contiguous.py | 51 ++++++++++++++ .../{contiguous.py => make_contiguous.py} | 43 ------------ .../tests/is_contiguous.py | 70 +++++++++++++++++++ .../{contiguous.py => make_contiguous.py} | 0 9 files changed, 144 insertions(+), 49 deletions(-) create mode 100644 docs/src/operations/reference/logic/is_contiguous.rst create mode 100644 docs/src/operations/reference/manipulation/make_contiguous.rst create mode 100644 python/metatensor-operations/metatensor/operations/is_contiguous.py rename python/metatensor-operations/metatensor/operations/{contiguous.py => make_contiguous.py} (57%) create mode 100644 python/metatensor-operations/tests/is_contiguous.py rename python/metatensor-operations/tests/{contiguous.py => make_contiguous.py} (100%) diff --git a/docs/src/operations/reference/logic/index.rst b/docs/src/operations/reference/logic/index.rst index a14a1c814..badd20b6b 100644 --- a/docs/src/operations/reference/logic/index.rst +++ b/docs/src/operations/reference/logic/index.rst @@ -7,3 +7,4 @@ Logic function allclose() equal() equal_metadata() + is_contiguous() diff --git a/docs/src/operations/reference/logic/is_contiguous.rst b/docs/src/operations/reference/logic/is_contiguous.rst new file mode 100644 index 000000000..abf6a59d9 --- /dev/null +++ b/docs/src/operations/reference/logic/is_contiguous.rst @@ -0,0 +1,7 @@ +is_contiguous +============= + +.. autofunction:: metatensor.is_contiguous + +.. autofunction:: metatensor.is_contiguous_block + diff --git a/docs/src/operations/reference/manipulation/index.rst b/docs/src/operations/reference/manipulation/index.rst index a53ef844d..21b9c2dd3 100644 --- a/docs/src/operations/reference/manipulation/index.rst +++ b/docs/src/operations/reference/manipulation/index.rst @@ -7,6 +7,7 @@ Manipulation operations detach() drop_blocks() join() + make_contiguous() manipulate dimension one_hot() remove_gradients() diff --git a/docs/src/operations/reference/manipulation/make_contiguous.rst b/docs/src/operations/reference/manipulation/make_contiguous.rst new file mode 100644 index 000000000..11589a0fc --- /dev/null +++ b/docs/src/operations/reference/manipulation/make_contiguous.rst @@ -0,0 +1,6 @@ +make_contiguous +=============== + +.. autofunction:: metatensor.make_contiguous + +.. autofunction:: metatensor.make_contiguous_block diff --git a/python/metatensor-operations/metatensor/operations/__init__.py b/python/metatensor-operations/metatensor/operations/__init__.py index c7adcc509..d49cd0706 100644 --- a/python/metatensor-operations/metatensor/operations/__init__.py +++ b/python/metatensor-operations/metatensor/operations/__init__.py @@ -15,12 +15,6 @@ allclose_raise, ) from .block_from_array import block_from_array # noqa -from .contiguous import ( # noqa: F401 - is_contiguous, - is_contiguous_block, - make_contiguous, - make_contiguous_block, -) from .detach import detach, detach_block # noqa from .divide import divide # noqa from .dot import dot # noqa @@ -33,8 +27,16 @@ equal_metadata_raise, equal_metadata_block_raise, ) +from .is_contiguous import ( # noqa: F401 + is_contiguous, + is_contiguous_block, +) from .join import join # noqa from .lstsq import lstsq # noqa +from .make_contiguous import ( # noqa: F401 + make_contiguous, + make_contiguous_block, +) from .manipulate_dimension import ( # noqa append_dimension, insert_dimension, diff --git a/python/metatensor-operations/metatensor/operations/is_contiguous.py b/python/metatensor-operations/metatensor/operations/is_contiguous.py new file mode 100644 index 000000000..b7e5b741b --- /dev/null +++ b/python/metatensor-operations/metatensor/operations/is_contiguous.py @@ -0,0 +1,51 @@ +from . import _dispatch +from ._backend import ( + Labels, + TensorBlock, + TensorMap, + torch_jit_is_scripting, + torch_jit_script, +) + + +@torch_jit_script +def is_contiguous_block(block: TensorBlock) -> bool: + """ + Checks whether the values array and gradients values arrays (if present) of an input + :py:class:`TensorBlock` are contiguous. + + Note that arrays of :py:class:`Labels` objects are not checked for contiguity. + + :param block: the input :py:class:`TensorBlock`. + + :return: bool, true if all values arrays contiguous, false otherwise. + """ + check_contiguous: bool = True + if not _dispatch.is_contiguous_array(block.values): + check_contiguous = False + + for _param, gradient in block.gradients(): + if not _dispatch.is_contiguous_array(gradient.values): + check_contiguous = False + + return check_contiguous + + +@torch_jit_script +def is_contiguous(tensor: TensorMap) -> bool: + """ + Checks whether all values arrays and gradients values arrays (if present) in all + :py:class:`TensorBlock` of an input :py:class:`TensorMap` are contiguous. + + Note that arrays of :py:class:`Labels` objects are not checked for contiguity. + + :param tensor: the input :py:class:`TensorMap`. + + :return: bool, true if all values arrays contiguous, false otherwise. + """ + check_contiguous: bool = True + for _key, block in tensor.items(): + if not is_contiguous_block(block): + check_contiguous = False + + return check_contiguous diff --git a/python/metatensor-operations/metatensor/operations/contiguous.py b/python/metatensor-operations/metatensor/operations/make_contiguous.py similarity index 57% rename from python/metatensor-operations/metatensor/operations/contiguous.py rename to python/metatensor-operations/metatensor/operations/make_contiguous.py index e9996d4e6..99613e004 100644 --- a/python/metatensor-operations/metatensor/operations/contiguous.py +++ b/python/metatensor-operations/metatensor/operations/make_contiguous.py @@ -8,49 +8,6 @@ ) -@torch_jit_script -def is_contiguous_block(block: TensorBlock) -> bool: - """ - Checks whether the values array and gradients values arrays (if present) of an input - :py:class:`TensorBlock` are contiguous. - - Note that arrays of :py:class:`Labels` objects are not checked for contiguity. - - :param block: the input :py:class:`TensorBlock`. - - :return: bool, true if all values arrays contiguous, false otherwise. - """ - check_contiguous: bool = True - if not _dispatch.is_contiguous_array(block.values): - check_contiguous = False - - for _param, gradient in block.gradients(): - if not is_contiguous_block(gradient): - check_contiguous = False - - return check_contiguous - - -@torch_jit_script -def is_contiguous(tensor: TensorMap) -> bool: - """ - Checks whether all values arrays and gradients values arrays (if present) in all - :py:class:`TensorBlock` of an input :py:class:`TensorMap` are contiguous. - - Note that arrays of :py:class:`Labels` objects are not checked for contiguity. - - :param tensor: the input :py:class:`TensorMap`. - - :return: bool, true if all values arrays contiguous, false otherwise. - """ - check_contiguous: bool = True - for _key, block in tensor.items(): - if not is_contiguous_block(block): - check_contiguous = False - - return check_contiguous - - @torch_jit_script def make_contiguous_block(block: TensorBlock) -> TensorBlock: """ diff --git a/python/metatensor-operations/tests/is_contiguous.py b/python/metatensor-operations/tests/is_contiguous.py new file mode 100644 index 000000000..d92d34735 --- /dev/null +++ b/python/metatensor-operations/tests/is_contiguous.py @@ -0,0 +1,70 @@ +import os + +import pytest + +import metatensor +from metatensor import TensorBlock, TensorMap + + +DATA_ROOT = os.path.join(os.path.dirname(__file__), "data") + + +@pytest.fixture +def tensor(): + """Loads a TensorMap from file for use in tests""" + return metatensor.load( + os.path.join(DATA_ROOT, "qm7-power-spectrum.npz"), + use_numpy=True, + ) + + +@pytest.fixture +def incontiguous_tensor(tensor) -> TensorMap: + """ + Make a TensorMap non-contiguous by reversing the order of the samples/properties + rows/columns in all values and gradient blocks. + """ + keys = tensor.keys + new_blocks = [] + + for _key, block in tensor.items(): + # Create a new TM with a non-contig array + new_block = _incontiguous_block(block) + new_blocks.append(new_block) + + return TensorMap(keys=keys, blocks=new_blocks) + + +def _incontiguous_block(block: TensorBlock) -> TensorBlock: + """ + Make a non-contiguous block by reversing the order in both the main value block and + the gradient block(s). + """ + new_vals = block.values.copy()[::-1, ::-1] + new_block = TensorBlock( + values=new_vals, + samples=block.samples, + components=block.components, + properties=block.properties, + ) + for param, gradient in block.gradients(): + new_grad = gradient.values.copy()[::-1, ::-1] + new_gradient = TensorBlock( + values=new_grad, + samples=gradient.samples, + components=gradient.components, + properties=gradient.properties, + ) + new_block.add_gradient(param, new_gradient) + + return new_block + + +def test_is_contiguous_block(tensor): + assert metatensor.is_contiguous_block(tensor.block(0)) + assert not metatensor.is_contiguous_block(_incontiguous_block(tensor.block(0))) + + +def test_is_contiguous(tensor, incontiguous_tensor): + assert metatensor.is_contiguous(tensor) + assert not metatensor.is_contiguous(incontiguous_tensor) \ No newline at end of file diff --git a/python/metatensor-operations/tests/contiguous.py b/python/metatensor-operations/tests/make_contiguous.py similarity index 100% rename from python/metatensor-operations/tests/contiguous.py rename to python/metatensor-operations/tests/make_contiguous.py From e2d5c019cfa98a9721fccc90096e24a207aa98d1 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Wed, 3 Jul 2024 16:52:25 +0200 Subject: [PATCH 6/7] Finish up the code and clean up tests --- ...{equal_metadata.rst => equal-metadata.rst} | 0 docs/src/operations/reference/logic/index.rst | 4 +- .../{is_contiguous.rst => is-contiguous.rst} | 0 ...ake_contiguous.rst => make-contiguous.rst} | 0 .../metatensor/operations/_dispatch.py | 9 ++-- .../metatensor/operations/is_contiguous.py | 31 +++++------- .../metatensor/operations/make_contiguous.py | 48 ++++++++----------- .../tests/is_contiguous.py | 15 +++--- .../tests/make_contiguous.py | 23 +++++---- .../tests/operations/contiguous.py | 13 +++-- 10 files changed, 68 insertions(+), 75 deletions(-) rename docs/src/operations/reference/logic/{equal_metadata.rst => equal-metadata.rst} (100%) rename docs/src/operations/reference/logic/{is_contiguous.rst => is-contiguous.rst} (100%) rename docs/src/operations/reference/manipulation/{make_contiguous.rst => make-contiguous.rst} (100%) diff --git a/docs/src/operations/reference/logic/equal_metadata.rst b/docs/src/operations/reference/logic/equal-metadata.rst similarity index 100% rename from docs/src/operations/reference/logic/equal_metadata.rst rename to docs/src/operations/reference/logic/equal-metadata.rst diff --git a/docs/src/operations/reference/logic/index.rst b/docs/src/operations/reference/logic/index.rst index badd20b6b..f125bc013 100644 --- a/docs/src/operations/reference/logic/index.rst +++ b/docs/src/operations/reference/logic/index.rst @@ -6,5 +6,5 @@ Logic function allclose() equal() - equal_metadata() - is_contiguous() + equal_metadata() + is_contiguous() diff --git a/docs/src/operations/reference/logic/is_contiguous.rst b/docs/src/operations/reference/logic/is-contiguous.rst similarity index 100% rename from docs/src/operations/reference/logic/is_contiguous.rst rename to docs/src/operations/reference/logic/is-contiguous.rst diff --git a/docs/src/operations/reference/manipulation/make_contiguous.rst b/docs/src/operations/reference/manipulation/make-contiguous.rst similarity index 100% rename from docs/src/operations/reference/manipulation/make_contiguous.rst rename to docs/src/operations/reference/manipulation/make-contiguous.rst diff --git a/python/metatensor-operations/metatensor/operations/_dispatch.py b/python/metatensor-operations/metatensor/operations/_dispatch.py index 88ede1920..e58f4bcc8 100644 --- a/python/metatensor-operations/metatensor/operations/_dispatch.py +++ b/python/metatensor-operations/metatensor/operations/_dispatch.py @@ -343,7 +343,7 @@ def concatenate(arrays: List[TorchTensor], axis: int): raise TypeError(UNKNOWN_ARRAY_TYPE) -def is_contiguous_array(array): +def is_contiguous(array): """ Checks if a given array is contiguous. @@ -360,11 +360,12 @@ def is_contiguous_array(array): raise TypeError(UNKNOWN_ARRAY_TYPE) -def make_contiguous_array(array): +def make_contiguous(array): """ Returns a contiguous array. - It is equivalent of np.ascontiguousarray(array) and tensor.contiguous(). In - the case of numpy, C order is used for consistency with torch. As such, only + + This is equivalent of ``np.ascontiguousarray(array)`` and ``tensor.contiguous()``. + In the case of numpy, C order is used for consistency with torch. As such, only C-contiguity is checked. """ if isinstance(array, TorchTensor): diff --git a/python/metatensor-operations/metatensor/operations/is_contiguous.py b/python/metatensor-operations/metatensor/operations/is_contiguous.py index b7e5b741b..77be4200b 100644 --- a/python/metatensor-operations/metatensor/operations/is_contiguous.py +++ b/python/metatensor-operations/metatensor/operations/is_contiguous.py @@ -1,11 +1,5 @@ from . import _dispatch -from ._backend import ( - Labels, - TensorBlock, - TensorMap, - torch_jit_is_scripting, - torch_jit_script, -) +from ._backend import TensorBlock, TensorMap, torch_jit_script @torch_jit_script @@ -20,15 +14,17 @@ def is_contiguous_block(block: TensorBlock) -> bool: :return: bool, true if all values arrays contiguous, false otherwise. """ - check_contiguous: bool = True - if not _dispatch.is_contiguous_array(block.values): - check_contiguous = False + if not _dispatch.is_contiguous(block.values): + return False - for _param, gradient in block.gradients(): - if not _dispatch.is_contiguous_array(gradient.values): - check_contiguous = False + for _parameter, gradient in block.gradients(): + if len(gradient.gradients_list()) != 0: + raise NotImplementedError("gradients of gradients are not supported") - return check_contiguous + if not _dispatch.is_contiguous(gradient.values): + return False + + return True @torch_jit_script @@ -43,9 +39,8 @@ def is_contiguous(tensor: TensorMap) -> bool: :return: bool, true if all values arrays contiguous, false otherwise. """ - check_contiguous: bool = True - for _key, block in tensor.items(): + for block in tensor.blocks(): if not is_contiguous_block(block): - check_contiguous = False + return False - return check_contiguous + return True diff --git a/python/metatensor-operations/metatensor/operations/make_contiguous.py b/python/metatensor-operations/metatensor/operations/make_contiguous.py index 99613e004..880057dda 100644 --- a/python/metatensor-operations/metatensor/operations/make_contiguous.py +++ b/python/metatensor-operations/metatensor/operations/make_contiguous.py @@ -1,57 +1,49 @@ +from typing import List + from . import _dispatch -from ._backend import ( - Labels, - TensorBlock, - TensorMap, - torch_jit_is_scripting, - torch_jit_script, -) +from ._backend import TensorBlock, TensorMap, torch_jit_script @torch_jit_script def make_contiguous_block(block: TensorBlock) -> TensorBlock: """ - Returns a new :py:class:`TensorBlock` where the values and gradient values (if - present) arrays are mades to be contiguous. + Returns a new :py:class:`TensorBlock` where the values and gradient (if present) + arrays are made to be contiguous. :param block: the input :py:class:`TensorBlock`. - - :return: a new :py:class:`TensorBlock` where the values and gradients arrays (if - present) are contiguous. """ - contiguous_block = TensorBlock( - values=_dispatch.make_contiguous_array(block.values.copy()), + new_block = TensorBlock( + values=_dispatch.make_contiguous(block.values), samples=block.samples, components=block.components, properties=block.properties, ) - for param, gradient in block.gradients(): + + for parameter, gradient in block.gradients(): + if len(gradient.gradients_list()) != 0: + raise NotImplementedError("gradients of gradients are not supported") + new_gradient = TensorBlock( - values=_dispatch.make_contiguous_array(gradient.values.copy()), + values=_dispatch.make_contiguous(gradient.values), samples=gradient.samples, components=gradient.components, properties=gradient.properties, ) - contiguous_block.add_gradient(param, new_gradient) + new_block.add_gradient(parameter, new_gradient) - return contiguous_block + return new_block @torch_jit_script def make_contiguous(tensor: TensorMap) -> TensorMap: """ Returns a new :py:class:`TensorMap` where all values and gradient values arrays are - mades to be contiguous. + made to be contiguous. :param tensor: the input :py:class:`TensorMap`. - - :return: a new :py:class:`TensorMap` with the same data and metadata as ``tensor`` - and contiguous values of ``tensor``. """ - keys: Labels = tensor.keys - contiguous_blocks: List[TensorBlock] = [] - for _key, block in tensor.items(): - contiguous_block = make_contiguous_block(block) - contiguous_blocks.append(contiguous_block) + new_blocks: List[TensorBlock] = [] + for block in tensor.blocks(): + new_blocks.append(make_contiguous_block(block)) - return TensorMap(keys=keys, blocks=contiguous_blocks) + return TensorMap(tensor.keys, new_blocks) diff --git a/python/metatensor-operations/tests/is_contiguous.py b/python/metatensor-operations/tests/is_contiguous.py index d92d34735..dab00409a 100644 --- a/python/metatensor-operations/tests/is_contiguous.py +++ b/python/metatensor-operations/tests/is_contiguous.py @@ -19,7 +19,7 @@ def tensor(): @pytest.fixture -def incontiguous_tensor(tensor) -> TensorMap: +def non_contiguous_tensor(tensor) -> TensorMap: """ Make a TensorMap non-contiguous by reversing the order of the samples/properties rows/columns in all values and gradient blocks. @@ -27,15 +27,14 @@ def incontiguous_tensor(tensor) -> TensorMap: keys = tensor.keys new_blocks = [] - for _key, block in tensor.items(): - # Create a new TM with a non-contig array - new_block = _incontiguous_block(block) + for block in tensor.blocks(): + new_block = _non_contiguous_block(block) new_blocks.append(new_block) return TensorMap(keys=keys, blocks=new_blocks) -def _incontiguous_block(block: TensorBlock) -> TensorBlock: +def _non_contiguous_block(block: TensorBlock) -> TensorBlock: """ Make a non-contiguous block by reversing the order in both the main value block and the gradient block(s). @@ -62,9 +61,9 @@ def _incontiguous_block(block: TensorBlock) -> TensorBlock: def test_is_contiguous_block(tensor): assert metatensor.is_contiguous_block(tensor.block(0)) - assert not metatensor.is_contiguous_block(_incontiguous_block(tensor.block(0))) + assert not metatensor.is_contiguous_block(_non_contiguous_block(tensor.block(0))) -def test_is_contiguous(tensor, incontiguous_tensor): +def test_is_contiguous(tensor, non_contiguous_tensor): assert metatensor.is_contiguous(tensor) - assert not metatensor.is_contiguous(incontiguous_tensor) \ No newline at end of file + assert not metatensor.is_contiguous(non_contiguous_tensor) diff --git a/python/metatensor-operations/tests/make_contiguous.py b/python/metatensor-operations/tests/make_contiguous.py index 2f66d8a13..5237f1a81 100644 --- a/python/metatensor-operations/tests/make_contiguous.py +++ b/python/metatensor-operations/tests/make_contiguous.py @@ -19,7 +19,7 @@ def tensor(): @pytest.fixture -def incontiguous_tensor(tensor) -> TensorMap: +def non_contiguous_tensor(tensor) -> TensorMap: """ Make a TensorMap non-contiguous by reversing the order of the samples/properties rows/columns in all values and gradient blocks. @@ -28,14 +28,13 @@ def incontiguous_tensor(tensor) -> TensorMap: new_blocks = [] for _key, block in tensor.items(): - # Create a new TM with a non-contig array - new_block = _incontiguous_block(block) + new_block = _non_contiguous_block(block) new_blocks.append(new_block) return TensorMap(keys=keys, blocks=new_blocks) -def _incontiguous_block(block: TensorBlock) -> TensorBlock: +def _non_contiguous_block(block: TensorBlock) -> TensorBlock: """ Make a non-contiguous block by reversing the order in both the main value block and the gradient block(s). @@ -60,13 +59,13 @@ def _incontiguous_block(block: TensorBlock) -> TensorBlock: return new_block -def test_is_contiguous_block(tensor): - assert not metatensor.is_contiguous_block(_incontiguous_block(tensor.block(0))) - assert metatensor.is_contiguous_block( - metatensor.make_contiguous_block(_incontiguous_block(tensor.block(0))) - ) +def test_make_contiguous_block(tensor): + block = _non_contiguous_block(tensor.block(0)) + + assert not metatensor.is_contiguous_block(block) + assert metatensor.is_contiguous_block(metatensor.make_contiguous_block(block)) -def test_is_contiguous(incontiguous_tensor): - assert not metatensor.is_contiguous(incontiguous_tensor) - assert metatensor.is_contiguous(metatensor.make_contiguous(incontiguous_tensor)) \ No newline at end of file +def test_make_contiguous(non_contiguous_tensor): + assert not metatensor.is_contiguous(non_contiguous_tensor) + assert metatensor.is_contiguous(metatensor.make_contiguous(non_contiguous_tensor)) diff --git a/python/metatensor-torch/tests/operations/contiguous.py b/python/metatensor-torch/tests/operations/contiguous.py index c6919a390..cd07099f2 100644 --- a/python/metatensor-torch/tests/operations/contiguous.py +++ b/python/metatensor-torch/tests/operations/contiguous.py @@ -5,6 +5,13 @@ import metatensor.torch -def test_is_contiguous(): - # TODO: write tests, used as a placeholder for now - assert True +def test_save_load(): + with io.BytesIO() as buffer: + torch.jit.save(metatensor.torch.is_contiguous, buffer) + buffer.seek(0) + torch.jit.load(buffer) + + with io.BytesIO() as buffer: + torch.jit.save(metatensor.torch.make_contiguous, buffer) + buffer.seek(0) + torch.jit.load(buffer) From 1313a85d49aeddc09440bb985183f55e71338880 Mon Sep 17 00:00:00 2001 From: Guillaume Fraux Date: Thu, 4 Jul 2024 15:49:10 +0200 Subject: [PATCH 7/7] Adress review --- .../manipulation/make-contiguous.rst | 3 ++ .../tests/is_contiguous.py | 32 +++++++++---------- .../tests/make_contiguous.py | 32 +++++++++---------- 3 files changed, 35 insertions(+), 32 deletions(-) diff --git a/docs/src/operations/reference/manipulation/make-contiguous.rst b/docs/src/operations/reference/manipulation/make-contiguous.rst index 11589a0fc..504247f15 100644 --- a/docs/src/operations/reference/manipulation/make-contiguous.rst +++ b/docs/src/operations/reference/manipulation/make-contiguous.rst @@ -1,6 +1,9 @@ make_contiguous =============== +These operations are equivalent of :py:func:`numpy.ascontiguousarray` in numpy +and :py:meth:`torch.Tensor.contiguous` in torch. + .. autofunction:: metatensor.make_contiguous .. autofunction:: metatensor.make_contiguous_block diff --git a/python/metatensor-operations/tests/is_contiguous.py b/python/metatensor-operations/tests/is_contiguous.py index dab00409a..258933812 100644 --- a/python/metatensor-operations/tests/is_contiguous.py +++ b/python/metatensor-operations/tests/is_contiguous.py @@ -18,22 +18,6 @@ def tensor(): ) -@pytest.fixture -def non_contiguous_tensor(tensor) -> TensorMap: - """ - Make a TensorMap non-contiguous by reversing the order of the samples/properties - rows/columns in all values and gradient blocks. - """ - keys = tensor.keys - new_blocks = [] - - for block in tensor.blocks(): - new_block = _non_contiguous_block(block) - new_blocks.append(new_block) - - return TensorMap(keys=keys, blocks=new_blocks) - - def _non_contiguous_block(block: TensorBlock) -> TensorBlock: """ Make a non-contiguous block by reversing the order in both the main value block and @@ -59,6 +43,22 @@ def _non_contiguous_block(block: TensorBlock) -> TensorBlock: return new_block +@pytest.fixture +def non_contiguous_tensor(tensor) -> TensorMap: + """ + Make a TensorMap non-contiguous by reversing the order of the samples/properties + rows/columns in all values and gradient blocks. + """ + keys = tensor.keys + new_blocks = [] + + for block in tensor.blocks(): + new_block = _non_contiguous_block(block) + new_blocks.append(new_block) + + return TensorMap(keys=keys, blocks=new_blocks) + + def test_is_contiguous_block(tensor): assert metatensor.is_contiguous_block(tensor.block(0)) assert not metatensor.is_contiguous_block(_non_contiguous_block(tensor.block(0))) diff --git a/python/metatensor-operations/tests/make_contiguous.py b/python/metatensor-operations/tests/make_contiguous.py index 5237f1a81..5af2ec3ae 100644 --- a/python/metatensor-operations/tests/make_contiguous.py +++ b/python/metatensor-operations/tests/make_contiguous.py @@ -18,22 +18,6 @@ def tensor(): ) -@pytest.fixture -def non_contiguous_tensor(tensor) -> TensorMap: - """ - Make a TensorMap non-contiguous by reversing the order of the samples/properties - rows/columns in all values and gradient blocks. - """ - keys = tensor.keys - new_blocks = [] - - for _key, block in tensor.items(): - new_block = _non_contiguous_block(block) - new_blocks.append(new_block) - - return TensorMap(keys=keys, blocks=new_blocks) - - def _non_contiguous_block(block: TensorBlock) -> TensorBlock: """ Make a non-contiguous block by reversing the order in both the main value block and @@ -59,6 +43,22 @@ def _non_contiguous_block(block: TensorBlock) -> TensorBlock: return new_block +@pytest.fixture +def non_contiguous_tensor(tensor) -> TensorMap: + """ + Make a TensorMap non-contiguous by reversing the order of the samples/properties + rows/columns in all values and gradient blocks. + """ + keys = tensor.keys + new_blocks = [] + + for _key, block in tensor.items(): + new_block = _non_contiguous_block(block) + new_blocks.append(new_block) + + return TensorMap(keys=keys, blocks=new_blocks) + + def test_make_contiguous_block(tensor): block = _non_contiguous_block(tensor.block(0))