{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Heat Basics" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## What is Heat for?\n", "---\n", "\n", "\n", "Straight from our [GitHub repository](https://github.com/helmholtz-analytics/heat):\n", "\n", "Heat builds on [PyTorch](https://pytorch.org/) and [mpi4py](https://mpi4py.readthedocs.io) to provide high-performance computing infrastructure for memory-intensive applications within the NumPy/SciPy ecosystem.\n", "\n", "\n", "With Heat you can:\n", "- port existing NumPy/SciPy code from single-CPU to multi-node clusters with minimal coding effort;\n", "- exploit the entire, cumulative RAM of your many nodes for memory-intensive operations and algorithms;\n", "- run your NumPy/SciPy code on GPUs (CUDA, ROCm, Apple MPS, coming up: Intel XPUs)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Why?\n", "\n", "- significant **scalability** with respect to task-parallel frameworks;\n", "- analysis of massive datasets without breaking them up in artificially independent chunks;\n", "- ease of use: script and test on your laptop, port straight to HPC cluster; \n", "- PyTorch-based: GPU support beyond the CUDA ecosystem." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", " \n", " \n", " \n", "
\n", "\n", "*(from left to right: Strong scaling on CPUs, weak scaling on CPUs, Weak scaling on GPUs)*" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Connecting to ipyparallel cluster\n", "\n", "We have started an `ipcluster` with 4 engines at the end of the [Setup notebook](0_setup/0_setup_local.ipynb).\n", "\n", "Let's start the interactive session with a look into the `heat` data object. But first, we need to import the `ipyparallel` client." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "4 engines found\n" ] } ], "source": [ "from ipyparallel import Client\n", "rc = Client(profile=\"default\")\n", "rc.ids\n", "\n", "if len(rc.ids) == 0:\n", " print(\"No engines found\")\n", "else:\n", " print(f\"{len(rc.ids)} engines found\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We will always start `heat` cells with the `%%px` magic command to execute the cell on all engines. However, the first section of this tutorial doesn't deal with distributed arrays. \n", "\n", "---" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## DNDarrays\n", "\n", "\n", "Similar to a NumPy `ndarray`, a Heat `dndarray` (we'll get to the `d` later) is a grid of values of a single (one particular) type. The number of dimensions is the number of axes of the array, while the shape of an array is a tuple of integers giving the number of elements of the array along each dimension. \n", "\n", "Heat emulates NumPy's API as closely as possible, allowing for the use of well-known **array creation functions**." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [] }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "6427fcd13db2489a878613869add8bfb", "version_major": 2, "version_minor": 0 }, "text/plain": [ "%px: 0%| | 0/4 [00:00 'DNDarray'\n", " Sum of array elements over a given axis. An array with the same shape as ``self.__array`` except\n", " for the specified axis which becomes one, e.g.\n", " ``a.shape=(1, 2, 3)`` => ``ht.ones((1, 2, 3)).sum(axis=1).shape=(1, 1, 3)``\n", " \n", " Parameters\n", " ----------\n", " a : DNDarray\n", " Input array.\n", " axis : None or int or Tuple[int,...], optional\n", " Axis along which a sum is performed. The default, ``axis=None``, will sum all of the\n", " elements of the input array. If ``axis`` is negative it counts from the last to the first\n", " axis. If ``axis`` is a tuple of ints, a sum is performed on all of the axes specified in the\n", " tuple instead of a single axis or all the axes as before.\n", " out : DNDarray, optional\n", " Alternative output array in which to place the result. It must have the same shape as the\n", " expected output, but the datatype of the output values will be cast if necessary.\n", " keepdims : bool, optional\n", " If this is set to ``True``, the axes which are reduced are left in the result as dimensions\n", " with size one. With this option, the result will broadcast correctly against the input\n", " array.\n", " \n", " Examples\n", " --------\n", " >>> ht.sum(ht.ones(2))\n", " DNDarray(2., dtype=ht.float32, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 3)))\n", " DNDarray(9., dtype=ht.float32, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 3)).astype(ht.int))\n", " DNDarray(9, dtype=ht.int64, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 2, 1)), axis=-3)\n", " DNDarray([[3.],\n", " [3.]], dtype=ht.float32, device=cpu:0, split=None)\n", "\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:0] Help on function sum in module heat.core.arithmetics:\n", "\n", "sum(a: 'DNDarray', axis: 'Union[int, Tuple[int, ...]]' = None, out: 'DNDarray' = None, keepdims: 'bool' = None) -> 'DNDarray'\n", " Sum of array elements over a given axis. An array with the same shape as ``self.__array`` except\n", " for the specified axis which becomes one, e.g.\n", " ``a.shape=(1, 2, 3)`` => ``ht.ones((1, 2, 3)).sum(axis=1).shape=(1, 1, 3)``\n", " \n", " Parameters\n", " ----------\n", " a : DNDarray\n", " Input array.\n", " axis : None or int or Tuple[int,...], optional\n", " Axis along which a sum is performed. The default, ``axis=None``, will sum all of the\n", " elements of the input array. If ``axis`` is negative it counts from the last to the first\n", " axis. If ``axis`` is a tuple of ints, a sum is performed on all of the axes specified in the\n", " tuple instead of a single axis or all the axes as before.\n", " out : DNDarray, optional\n", " Alternative output array in which to place the result. It must have the same shape as the\n", " expected output, but the datatype of the output values will be cast if necessary.\n", " keepdims : bool, optional\n", " If this is set to ``True``, the axes which are reduced are left in the result as dimensions\n", " with size one. With this option, the result will broadcast correctly against the input\n", " array.\n", " \n", " Examples\n", " --------\n", " >>> ht.sum(ht.ones(2))\n", " DNDarray(2., dtype=ht.float32, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 3)))\n", " DNDarray(9., dtype=ht.float32, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 3)).astype(ht.int))\n", " DNDarray(9, dtype=ht.int64, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 2, 1)), axis=-3)\n", " DNDarray([[3.],\n", " [3.]], dtype=ht.float32, device=cpu:0, split=None)\n", "\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:3] Help on function sum in module heat.core.arithmetics:\n", "\n", "sum(a: 'DNDarray', axis: 'Union[int, Tuple[int, ...]]' = None, out: 'DNDarray' = None, keepdims: 'bool' = None) -> 'DNDarray'\n", " Sum of array elements over a given axis. An array with the same shape as ``self.__array`` except\n", " for the specified axis which becomes one, e.g.\n", " ``a.shape=(1, 2, 3)`` => ``ht.ones((1, 2, 3)).sum(axis=1).shape=(1, 1, 3)``\n", " \n", " Parameters\n", " ----------\n", " a : DNDarray\n", " Input array.\n", " axis : None or int or Tuple[int,...], optional\n", " Axis along which a sum is performed. The default, ``axis=None``, will sum all of the\n", " elements of the input array. If ``axis`` is negative it counts from the last to the first\n", " axis. If ``axis`` is a tuple of ints, a sum is performed on all of the axes specified in the\n", " tuple instead of a single axis or all the axes as before.\n", " out : DNDarray, optional\n", " Alternative output array in which to place the result. It must have the same shape as the\n", " expected output, but the datatype of the output values will be cast if necessary.\n", " keepdims : bool, optional\n", " If this is set to ``True``, the axes which are reduced are left in the result as dimensions\n", " with size one. With this option, the result will broadcast correctly against the input\n", " array.\n", " \n", " Examples\n", " --------\n", " >>> ht.sum(ht.ones(2))\n", " DNDarray(2., dtype=ht.float32, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 3)))\n", " DNDarray(9., dtype=ht.float32, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 3)).astype(ht.int))\n", " DNDarray(9, dtype=ht.int64, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 2, 1)), axis=-3)\n", " DNDarray([[3.],\n", " [3.]], dtype=ht.float32, device=cpu:0, split=None)\n", "\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:1] Help on function sum in module heat.core.arithmetics:\n", "\n", "sum(a: 'DNDarray', axis: 'Union[int, Tuple[int, ...]]' = None, out: 'DNDarray' = None, keepdims: 'bool' = None) -> 'DNDarray'\n", " Sum of array elements over a given axis. An array with the same shape as ``self.__array`` except\n", " for the specified axis which becomes one, e.g.\n", " ``a.shape=(1, 2, 3)`` => ``ht.ones((1, 2, 3)).sum(axis=1).shape=(1, 1, 3)``\n", " \n", " Parameters\n", " ----------\n", " a : DNDarray\n", " Input array.\n", " axis : None or int or Tuple[int,...], optional\n", " Axis along which a sum is performed. The default, ``axis=None``, will sum all of the\n", " elements of the input array. If ``axis`` is negative it counts from the last to the first\n", " axis. If ``axis`` is a tuple of ints, a sum is performed on all of the axes specified in the\n", " tuple instead of a single axis or all the axes as before.\n", " out : DNDarray, optional\n", " Alternative output array in which to place the result. It must have the same shape as the\n", " expected output, but the datatype of the output values will be cast if necessary.\n", " keepdims : bool, optional\n", " If this is set to ``True``, the axes which are reduced are left in the result as dimensions\n", " with size one. With this option, the result will broadcast correctly against the input\n", " array.\n", " \n", " Examples\n", " --------\n", " >>> ht.sum(ht.ones(2))\n", " DNDarray(2., dtype=ht.float32, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 3)))\n", " DNDarray(9., dtype=ht.float32, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 3)).astype(ht.int))\n", " DNDarray(9, dtype=ht.int64, device=cpu:0, split=None)\n", " >>> ht.sum(ht.ones((3, 2, 1)), axis=-3)\n", " DNDarray([[3.],\n", " [3.]], dtype=ht.float32, device=cpu:0, split=None)\n", "\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "%%px \n", "help(ht.sum)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Parallel Processing\n", "\n", "Heat's actual power lies in the possibility to exploit the processing performance of modern accelerator hardware (GPUs) as well as distributed (high-performance) cluster systems. All operations executed on CPUs are, to a large extent, vectorized (AVX) and thread-parallelized (OpenMP). Heat builds on PyTorch, so it supports GPU acceleration on Nvidia and AMD GPUs. \n", "\n", "For distributed computations, your system or laptop needs to have Message Passing Interface (MPI) installed. For GPU computations, your system needs to have one or more suitable GPUs and (MPI-aware) CUDA/ROCm ecosystem.\n", "\n", "**NOTE:** The GPU examples below will only properly execute on a computer with a GPU. Make sure to either start the notebook on an appropriate machine or copy and paste the examples into a script and execute it on a suitable device." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### GPUs\n", "\n", "Heat's array creation functions all support an additional parameter that places the data on a specific device. By default, the CPU is selected, but it is also possible to directly allocate the data on a GPU." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "The following cells will only work if you have a GPU available.\n", "\n", "
" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "27a1a149c68b43e8a74fdcae79de2611", "version_major": 2, "version_minor": 0 }, "text/plain": [ "%px: 0%| | 0/4 [00:00\u001b[0;34m\u001b[0m\n", "\u001b[1;32m 1\u001b[0m \u001b[0ma\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mht\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfull\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'cpu'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 2\u001b[0m \u001b[0mb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mht\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mones\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'gpu'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m----> 3\u001b[0;31m \u001b[0ma\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\n", "\u001b[0;32m~/heat/heat/core/arithmetics.py\u001b[0m in \u001b[0;36m_add\u001b[0;34m(self, other)\u001b[0m\n", "\u001b[1;32m 122\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_add\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 123\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m--> 124\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0madd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 125\u001b[0m \u001b[0;32mexcept\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 126\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mNotImplemented\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;32m~/heat/heat/core/arithmetics.py\u001b[0m in \u001b[0;36madd\u001b[0;34m(t1, t2, out, where)\u001b[0m\n", "\u001b[1;32m 117\u001b[0m [5., 6.]], dtype=ht.float32, device=cpu:0, split=None)\n", "\u001b[1;32m 118\u001b[0m \"\"\"\n", "\u001b[0;32m--> 119\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0m_operations\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__binary_op\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwhere\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 120\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 121\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;32m~/heat/heat/core/_operations.py\u001b[0m in \u001b[0;36m__binary_op\u001b[0;34m(operation, t1, t2, out, where, fn_kwargs)\u001b[0m\n", "\u001b[1;32m 202\u001b[0m \u001b[0mpromoted_type\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat32\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 203\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m--> 204\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moperation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlarray\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpromoted_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlarray\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpromoted_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mfn_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 205\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 206\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mout\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mwhere\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;31mRuntimeError\u001b[0m: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cpu!\n", "[2:execute]\n", "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n", "\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)\n", "\u001b[0;32m/tmp/ipykernel_2337974/396898532.py\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n", "\u001b[1;32m 1\u001b[0m \u001b[0ma\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mht\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfull\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'cpu'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 2\u001b[0m \u001b[0mb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mht\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mones\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'gpu'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m----> 3\u001b[0;31m \u001b[0ma\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\n", "\u001b[0;32m~/heat/heat/core/arithmetics.py\u001b[0m in \u001b[0;36m_add\u001b[0;34m(self, other)\u001b[0m\n", "\u001b[1;32m 122\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_add\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 123\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m--> 124\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0madd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 125\u001b[0m \u001b[0;32mexcept\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 126\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mNotImplemented\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;32m~/heat/heat/core/arithmetics.py\u001b[0m in \u001b[0;36madd\u001b[0;34m(t1, t2, out, where)\u001b[0m\n", "\u001b[1;32m 117\u001b[0m [5., 6.]], dtype=ht.float32, device=cpu:0, split=None)\n", "\u001b[1;32m 118\u001b[0m \"\"\"\n", "\u001b[0;32m--> 119\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0m_operations\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__binary_op\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwhere\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 120\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 121\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;32m~/heat/heat/core/_operations.py\u001b[0m in \u001b[0;36m__binary_op\u001b[0;34m(operation, t1, t2, out, where, fn_kwargs)\u001b[0m\n", "\u001b[1;32m 202\u001b[0m \u001b[0mpromoted_type\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat32\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 203\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m--> 204\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moperation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlarray\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpromoted_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlarray\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpromoted_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mfn_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 205\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 206\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mout\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mwhere\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;31mRuntimeError\u001b[0m: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!\n", "[0:execute]\n", "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n", "\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)\n", "\u001b[0;32m/tmp/ipykernel_2337972/396898532.py\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n", "\u001b[1;32m 1\u001b[0m \u001b[0ma\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mht\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfull\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'cpu'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 2\u001b[0m \u001b[0mb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mht\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mones\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'gpu'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m----> 3\u001b[0;31m \u001b[0ma\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\n", "\u001b[0;32m~/heat/heat/core/arithmetics.py\u001b[0m in \u001b[0;36m_add\u001b[0;34m(self, other)\u001b[0m\n", "\u001b[1;32m 122\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_add\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 123\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m--> 124\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0madd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 125\u001b[0m \u001b[0;32mexcept\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 126\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mNotImplemented\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;32m~/heat/heat/core/arithmetics.py\u001b[0m in \u001b[0;36madd\u001b[0;34m(t1, t2, out, where)\u001b[0m\n", "\u001b[1;32m 117\u001b[0m [5., 6.]], dtype=ht.float32, device=cpu:0, split=None)\n", "\u001b[1;32m 118\u001b[0m \"\"\"\n", "\u001b[0;32m--> 119\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0m_operations\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__binary_op\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwhere\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 120\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 121\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;32m~/heat/heat/core/_operations.py\u001b[0m in \u001b[0;36m__binary_op\u001b[0;34m(operation, t1, t2, out, where, fn_kwargs)\u001b[0m\n", "\u001b[1;32m 202\u001b[0m \u001b[0mpromoted_type\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat32\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 203\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m--> 204\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moperation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlarray\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpromoted_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlarray\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpromoted_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mfn_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 205\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 206\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mout\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mwhere\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;31mRuntimeError\u001b[0m: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!\n", "[3:execute]\n", "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n", "\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)\n", "\u001b[0;32m/tmp/ipykernel_2337975/396898532.py\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n", "\u001b[1;32m 1\u001b[0m \u001b[0ma\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mht\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfull\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'cpu'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 2\u001b[0m \u001b[0mb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mht\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mones\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'gpu'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m----> 3\u001b[0;31m \u001b[0ma\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\n", "\u001b[0;32m~/heat/heat/core/arithmetics.py\u001b[0m in \u001b[0;36m_add\u001b[0;34m(self, other)\u001b[0m\n", "\u001b[1;32m 122\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_add\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 123\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m--> 124\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0madd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 125\u001b[0m \u001b[0;32mexcept\u001b[0m \u001b[0mTypeError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 126\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mNotImplemented\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;32m~/heat/heat/core/arithmetics.py\u001b[0m in \u001b[0;36madd\u001b[0;34m(t1, t2, out, where)\u001b[0m\n", "\u001b[1;32m 117\u001b[0m [5., 6.]], dtype=ht.float32, device=cpu:0, split=None)\n", "\u001b[1;32m 118\u001b[0m \"\"\"\n", "\u001b[0;32m--> 119\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0m_operations\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__binary_op\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwhere\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 120\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 121\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;32m~/heat/heat/core/_operations.py\u001b[0m in \u001b[0;36m__binary_op\u001b[0;34m(operation, t1, t2, out, where, fn_kwargs)\u001b[0m\n", "\u001b[1;32m 202\u001b[0m \u001b[0mpromoted_type\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat32\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 203\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0;32m--> 204\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moperation\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlarray\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpromoted_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlarray\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpromoted_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mfn_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[0m\u001b[1;32m 205\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", "\u001b[1;32m 206\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mout\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mwhere\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", "\n", "\u001b[0;31mRuntimeError\u001b[0m: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cpu!\n" ] }, { "ename": "AlreadyDisplayedError", "evalue": "4 errors", "output_type": "error", "traceback": [ "4 errors" ] } ], "source": [ "%%px \n", "a = ht.full((3, 4,), 4, device='cpu')\n", "b = ht.ones((3, 4,), device='gpu')\n", "a + b" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "It is possible to explicitly move an array from one device to the other and back to avoid this error." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "80de6265-3d236a5e1eb6ffa587e57739", "error": null, "execute_input": "a = ht.full((3, 4,), 4, device='gpu')\na.cpu()\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 24, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:22.834794Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "44afd95b-fbaa3f38705a9e054856d875", "error": null, "execute_input": "a = ht.full((3, 4,), 4, device='gpu')\na.cpu()\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 24, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:22.834910Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:24]: \u001b[0m\n", "DNDarray([[4., 4., 4., 4.],\n", " [4., 4., 4., 4.],\n", " [4., 4., 4., 4.]], dtype=ht.float32, device=cpu:0, split=None)" ] }, "metadata": { "after": [], "completed": "2025-11-21T13:52:22.844301Z", "data": {}, "engine_id": 0, "engine_uuid": "b6abf056-e357819abb2ce1db79af7a26", "error": null, "execute_input": "a = ht.full((3, 4,), 4, device='gpu')\na.cpu()\n", "execute_result": { "data": { "text/plain": "DNDarray([[4., 4., 4., 4.],\n [4., 4., 4., 4.],\n [4., 4., 4., 4.]], dtype=ht.float32, device=cpu:0, split=None)" }, "execution_count": 24, "metadata": {} }, "follow": [], "is_broadcast": false, "is_coalescing": false, "msg_id": "37b4a2a2-1abf521f7da7e7e369dafe73_2339267_89", "outputs": [], "received": "2025-11-21T13:52:22.845849Z", "started": "2025-11-21T13:52:22.838672Z", "status": "ok", "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:22.834521Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": [], "completed": "2025-11-21T13:52:22.845141Z", "data": {}, "engine_id": 2, "engine_uuid": "ab9782c4-a0794501cf2af7341aeff65d", "error": null, "execute_input": "a = ht.full((3, 4,), 4, device='gpu')\na.cpu()\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 24, "metadata": {} }, "follow": [], "is_broadcast": false, "is_coalescing": false, "msg_id": "37b4a2a2-1abf521f7da7e7e369dafe73_2339267_91", "outputs": [], "received": "2025-11-21T13:52:22.846939Z", "started": "2025-11-21T13:52:22.839905Z", "status": "ok", "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:22.834819Z" }, "output_type": "display_data" } ], "source": [ "%%px \n", "a = ht.full((3, 4,), 4, device='gpu')\n", "a.cpu()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We'll put our multi-GPU setup to the test in the next section." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Distributed Computing\n", "\n", "Heat is also able to make use of distributed processing capabilities such as those in high-performance cluster systems. For this, Heat exploits the fact that the operations performed on a multi-dimensional array are usually identical for all data items. Hence, a data-parallel processing strategy can be chosen, where the total number of data items is equally divided among all processing nodes. An operation is then performed individually on the local data chunks and, if necessary, communicates partial results behind the scenes. A Heat array assumes the role of a virtual overlay of the local chunks and realizes and coordinates the computations - see the figure below for a visual representation of this concept.\n", "\n", "\n", "\n", "The chunks are always split along a singular dimension (i.e. 1-D domain decomposition) of the array. You can specify this in Heat by using the `split` parameter. This parameter is present in all relevant functions, such as array creation (`zeros(), ones(), ...`) or I/O (`load()`) functions. \n", "\n", "\n", "\n", "\n", "Examples are provided below. The result of an operation on DNDarays will in most cases preserve the `split` or distribution axis of the respective operands. However, in some cases the split axis might change. For example, a transpose of a Heat array will equally transpose the split axis. Furthermore, a reduction operation, e.g. `sum()`, that is performed across the split axis, might remove data partitions entirely. The respective function behaviors can be found in Heat's documentation.\n", "\n", "You may also modify the data partitioning of a Heat array by using the `resplit()` function. This allows you to repartition the data as you so choose. Please note, that this should be used sparingly and for small data amounts only, as it entails significant data copying across the network. Finally, a Heat array without any split, i.e. `split=None` (default), will result in redundant copies of data on each computation node.\n", "\n", "On a technical level, Heat follows the so-called [Bulk Synchronous Parallel (BSP)](https://en.wikipedia.org/wiki/Bulk_synchronous_parallel) processing model. For the network communication, Heat utilizes the [Message Passing Interface (MPI)](https://computing.llnl.gov/tutorials/mpi/), a *de facto* standard on modern high-performance computing systems. It is also possible to use MPI on your laptop or desktop computer. Respective software packages are available for all major operating systems. In order to run a Heat script, you need to start it slightly differently than you are probably used to. This\n", "\n", "```bash\n", "python ./my_script.py\n", "```\n", "\n", "becomes this instead:\n", "\n", "```bash\n", "mpirun -n python ./my_script.py\n", "```\n", "On an HPC cluster you'll of course use SBATCH or similar.\n", "\n", "\n", "Let's see some examples of working with distributed Heat arrays:" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In the following examples, we'll recreate the array shown in the figure, a 3-dimensional DNDarray of integers ranging from 0 to 59 (5 matrices of size (4,3)). " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "\u001b[0;31mOut[0:25]: \u001b[0m\n", "DNDarray([[[ 0, 1, 2],\n", " [ 3, 4, 5],\n", " [ 6, 7, 8],\n", " [ 9, 10, 11]],\n", "\n", " [[12, 13, 14],\n", " [15, 16, 17],\n", " [18, 19, 20],\n", " [21, 22, 23]],\n", "\n", " [[24, 25, 26],\n", " [27, 28, 29],\n", " [30, 31, 32],\n", " [33, 34, 35]],\n", "\n", " [[36, 37, 38],\n", " [39, 40, 41],\n", " [42, 43, 44],\n", " [45, 46, 47]],\n", "\n", " [[48, 49, 50],\n", " [51, 52, 53],\n", " [54, 55, 56],\n", " [57, 58, 59]]], dtype=ht.int32, device=cpu:0, split=None)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "b6abf056-e357819abb2ce1db79af7a26", "error": null, "execute_input": "import heat as ht\ndndarray = ht.arange(60).reshape(5,4,3)\ndndarray\n", "execute_result": { "data": { "text/plain": "DNDarray([[[ 0, 1, 2],\n [ 3, 4, 5],\n [ 6, 7, 8],\n [ 9, 10, 11]],\n\n [[12, 13, 14],\n [15, 16, 17],\n [18, 19, 20],\n [21, 22, 23]],\n\n [[24, 25, 26],\n [27, 28, 29],\n [30, 31, 32],\n [33, 34, 35]],\n\n [[36, 37, 38],\n [39, 40, 41],\n [42, 43, 44],\n [45, 46, 47]],\n\n [[48, 49, 50],\n [51, 52, 53],\n [54, 55, 56],\n [57, 58, 59]]], dtype=ht.int32, device=cpu:0, split=None)" }, "execution_count": 25, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:27.184903Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "ab9782c4-a0794501cf2af7341aeff65d", "error": null, "execute_input": "import heat as ht\ndndarray = ht.arange(60).reshape(5,4,3)\ndndarray\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 25, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:27.185473Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "80de6265-3d236a5e1eb6ffa587e57739", "error": null, "execute_input": "import heat as ht\ndndarray = ht.arange(60).reshape(5,4,3)\ndndarray\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 25, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:27.185317Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": [], "completed": "2025-11-21T13:52:27.194775Z", "data": {}, "engine_id": 3, "engine_uuid": "44afd95b-fbaa3f38705a9e054856d875", "error": null, "execute_input": "import heat as ht\ndndarray = ht.arange(60).reshape(5,4,3)\ndndarray\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 25, "metadata": {} }, "follow": [], "is_broadcast": false, "is_coalescing": false, "msg_id": "37b4a2a2-1abf521f7da7e7e369dafe73_2339267_96", "outputs": [], "received": "2025-11-21T13:52:27.196603Z", "started": "2025-11-21T13:52:27.190025Z", "status": "ok", "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:27.185506Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "import heat as ht\n", "dndarray = ht.arange(60).reshape(5,4,3)\n", "dndarray" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Notice the additional metadata printed with the DNDarray. With respect to a numpy ndarray, the DNDarray has additional information on the device (in this case, the CPU) and the `split` axis. In the example above, the split axis is `None`, meaning that the DNDarray is not distributed and each MPI process has a full copy of the data.\n", "\n", "Let's experiment with a distributed DNDarray: we'll split the same DNDarray as above, but distributed along the major axis." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "f7196622-1a114180b3ed026d0a484a97", "error": null, "execute_input": "dndarray = ht.arange(60, split=0).reshape(5,4,3)\ndndarray\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 2, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:12:09.011859Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "f34f97e0-8017bfe4727b64fead8923c2", "error": null, "execute_input": "dndarray = ht.arange(60, split=0).reshape(5,4,3)\ndndarray\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 2, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:12:09.011620Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:2]: \u001b[0m\n", "DNDarray([[[ 0, 1, 2],\n", " [ 3, 4, 5],\n", " [ 6, 7, 8],\n", " [ 9, 10, 11]],\n", "\n", " [[12, 13, 14],\n", " [15, 16, 17],\n", " [18, 19, 20],\n", " [21, 22, 23]],\n", "\n", " [[24, 25, 26],\n", " [27, 28, 29],\n", " [30, 31, 32],\n", " [33, 34, 35]],\n", "\n", " [[36, 37, 38],\n", " [39, 40, 41],\n", " [42, 43, 44],\n", " [45, 46, 47]],\n", "\n", " [[48, 49, 50],\n", " [51, 52, 53],\n", " [54, 55, 56],\n", " [57, 58, 59]]], dtype=ht.int32, device=cpu:0, split=0)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "dc94d760-c1cc422f82bcde0f12e02d04", "error": null, "execute_input": "dndarray = ht.arange(60, split=0).reshape(5,4,3)\ndndarray\n", "execute_result": { "data": { "text/plain": "DNDarray([[[ 0, 1, 2],\n [ 3, 4, 5],\n [ 6, 7, 8],\n [ 9, 10, 11]],\n\n [[12, 13, 14],\n [15, 16, 17],\n [18, 19, 20],\n [21, 22, 23]],\n\n [[24, 25, 26],\n [27, 28, 29],\n [30, 31, 32],\n [33, 34, 35]],\n\n [[36, 37, 38],\n [39, 40, 41],\n [42, 43, 44],\n [45, 46, 47]],\n\n [[48, 49, 50],\n [51, 52, 53],\n [54, 55, 56],\n [57, 58, 59]]], dtype=ht.int32, device=cpu:0, split=0)" }, "execution_count": 2, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:12:09.011513Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "5112c8cc-d495c807bbe7af0a6d4ab1ea", "error": null, "execute_input": "dndarray = ht.arange(60, split=0).reshape(5,4,3)\ndndarray\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 2, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T14:12:09.011676Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "dndarray = ht.arange(60, split=0).reshape(5,4,3)\n", "dndarray" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The `split` axis is now 0, meaning that the DNDarray is distributed along the first axis. Each MPI process has a slice of the data along the first axis. In order to see the data on each process, we can print the \"local array\" via the `larray` attribute." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "\u001b[0;31mOut[2:27]: \u001b[0m\n", "tensor([[[36, 37, 38],\n", " [39, 40, 41],\n", " [42, 43, 44],\n", " [45, 46, 47]]], dtype=torch.int32)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "ab9782c4-a0794501cf2af7341aeff65d", "error": null, "execute_input": "dndarray.larray\n", "execute_result": { "data": { "text/plain": "tensor([[[36, 37, 38],\n [39, 40, 41],\n [42, 43, 44],\n [45, 46, 47]]], dtype=torch.int32)" }, "execution_count": 27, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:45.523997Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[1:27]: \u001b[0m\n", "tensor([[[24, 25, 26],\n", " [27, 28, 29],\n", " [30, 31, 32],\n", " [33, 34, 35]]], dtype=torch.int32)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "80de6265-3d236a5e1eb6ffa587e57739", "error": null, "execute_input": "dndarray.larray\n", "execute_result": { "data": { "text/plain": "tensor([[[24, 25, 26],\n [27, 28, 29],\n [30, 31, 32],\n [33, 34, 35]]], dtype=torch.int32)" }, "execution_count": 27, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:45.523921Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:27]: \u001b[0m\n", "tensor([[[ 0, 1, 2],\n", " [ 3, 4, 5],\n", " [ 6, 7, 8],\n", " [ 9, 10, 11]],\n", "\n", " [[12, 13, 14],\n", " [15, 16, 17],\n", " [18, 19, 20],\n", " [21, 22, 23]]], dtype=torch.int32)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "b6abf056-e357819abb2ce1db79af7a26", "error": null, "execute_input": "dndarray.larray\n", "execute_result": { "data": { "text/plain": "tensor([[[ 0, 1, 2],\n [ 3, 4, 5],\n [ 6, 7, 8],\n [ 9, 10, 11]],\n\n [[12, 13, 14],\n [15, 16, 17],\n [18, 19, 20],\n [21, 22, 23]]], dtype=torch.int32)" }, "execution_count": 27, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:45.523823Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[3:27]: \u001b[0m\n", "tensor([[[48, 49, 50],\n", " [51, 52, 53],\n", " [54, 55, 56],\n", " [57, 58, 59]]], dtype=torch.int32)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "44afd95b-fbaa3f38705a9e054856d875", "error": null, "execute_input": "dndarray.larray\n", "execute_result": { "data": { "text/plain": "tensor([[[48, 49, 50],\n [51, 52, 53],\n [54, 55, 56],\n [57, 58, 59]]], dtype=torch.int32)" }, "execution_count": 27, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:45.524052Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "dndarray.larray" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Note that the `larray` is a `torch.Tensor` object. This is the underlying tensor that holds the data. The `dndarray` object is an MPI-aware wrapper around these process-local tensors, providing memory-distributed functionality and information." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The DNDarray can be distributed along any axis. Modify the `split` attribute when creating the DNDarray in the cell above, to distribute it along a different axis, and see how the `larray`s change. You'll notice that the distributed arrays are always load-balanced, meaning that the data is distributed as evenly as possible across the MPI processes." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The `DNDarray` object has a number of methods and attributes that are useful for distributed computing. In particular, it keeps track of its global and local (on a given process) shape through distributed operations and array manipulations. The DNDarray is also associated to a `comm` object, the MPI communicator.\n", "\n", "(In MPI, the *communicator* is a group of processes that can communicate with each other. The `comm` object is a `MPI.COMM_WORLD` communicator, which is the default communicator that includes all the processes. The `comm` object is used to perform collective operations, such as reductions, scatter, gather, and broadcast. The `comm` object is also used to perform point-to-point communication between processes.)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[stdout:1] Global shape of the dndarray: (5, 4, 3)\n", "On rank 1/4, local shape of the dndarray: (1, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:0] Global shape of the dndarray: (5, 4, 3)\n", "On rank 0/4, local shape of the dndarray: (2, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:2] Global shape of the dndarray: (5, 4, 3)\n", "On rank 2/4, local shape of the dndarray: (1, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:3] Global shape of the dndarray: (5, 4, 3)\n", "On rank 3/4, local shape of the dndarray: (1, 4, 3)\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "%%px\n", "print(f\"Global shape of the dndarray: {dndarray.shape}\")\n", "print(f\"On rank {dndarray.comm.rank}/{dndarray.comm.size}, local shape of the dndarray: {dndarray.lshape}\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can perform a vast number of operations on DNDarrays distributed over multi-node and/or multi-GPU resources. Check out our [Numpy coverage tables](https://github.com/helmholtz-analytics/heat/blob/main/coverage_tables.md) to see what operations are already supported. " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "ab9782c4-a0794501cf2af7341aeff65d", "error": null, "execute_input": "# transpose \ndndarray.T\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 29, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:54.418237Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "80de6265-3d236a5e1eb6ffa587e57739", "error": null, "execute_input": "# transpose \ndndarray.T\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 29, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:54.418162Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "44afd95b-fbaa3f38705a9e054856d875", "error": null, "execute_input": "# transpose \ndndarray.T\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 29, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:54.418303Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:29]: \u001b[0m\n", "DNDarray([[[ 0, 12, 24, 36, 48],\n", " [ 3, 15, 27, 39, 51],\n", " [ 6, 18, 30, 42, 54],\n", " [ 9, 21, 33, 45, 57]],\n", "\n", " [[ 1, 13, 25, 37, 49],\n", " [ 4, 16, 28, 40, 52],\n", " [ 7, 19, 31, 43, 55],\n", " [10, 22, 34, 46, 58]],\n", "\n", " [[ 2, 14, 26, 38, 50],\n", " [ 5, 17, 29, 41, 53],\n", " [ 8, 20, 32, 44, 56],\n", " [11, 23, 35, 47, 59]]], dtype=ht.int32, device=cpu:0, split=2)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "b6abf056-e357819abb2ce1db79af7a26", "error": null, "execute_input": "# transpose \ndndarray.T\n", "execute_result": { "data": { "text/plain": "DNDarray([[[ 0, 12, 24, 36, 48],\n [ 3, 15, 27, 39, 51],\n [ 6, 18, 30, 42, 54],\n [ 9, 21, 33, 45, 57]],\n\n [[ 1, 13, 25, 37, 49],\n [ 4, 16, 28, 40, 52],\n [ 7, 19, 31, 43, 55],\n [10, 22, 34, 46, 58]],\n\n [[ 2, 14, 26, 38, 50],\n [ 5, 17, 29, 41, 53],\n [ 8, 20, 32, 44, 56],\n [11, 23, 35, 47, 59]]], dtype=ht.int32, device=cpu:0, split=2)" }, "execution_count": 29, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:52:54.417998Z" }, "output_type": "display_data" } ], "source": [ "%%px \n", "# transpose \n", "dndarray.T\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[stdout:0] The slowest run took 32.65 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "200 µs ± 276 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:1] The slowest run took 22.67 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "169 µs ± 220 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:2] The slowest run took 37.03 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "194 µs ± 287 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:3] The slowest run took 41.47 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "192 µs ± 320 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "%%px\n", "# reduction operation along the distribution axis\n", "%timeit -n 1 dndarray.sum(axis=0)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[stdout:0] The slowest run took 30.48 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "136 µs ± 162 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:1] The slowest run took 27.16 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "193 µs ± 219 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:2] The slowest run took 18.36 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "67.1 µs ± 89.3 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:3] The slowest run took 18.05 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "61 µs ± 87.6 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "%%px \n", "# reduction operation along non-distribution axis: no communication required\n", "%timeit -n 1 dndarray.sum(axis=1)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Operations between tensors with equal split or no split are fully parallelizable and therefore very fast." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "80de6265-3d236a5e1eb6ffa587e57739", "error": null, "execute_input": "other_dndarray = ht.arange(60,120, split=0).reshape(5,4,3) # distributed reshape\n\n# element-wise multiplication\ndndarray * other_dndarray\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 32, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:53:09.351713Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "ab9782c4-a0794501cf2af7341aeff65d", "error": null, "execute_input": "other_dndarray = ht.arange(60,120, split=0).reshape(5,4,3) # distributed reshape\n\n# element-wise multiplication\ndndarray * other_dndarray\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 32, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:53:09.351803Z" }, "output_type": "display_data" }, { "data": { "text/plain": [ "\u001b[0;31mOut[0:32]: \u001b[0m\n", "DNDarray([[[ 0, 61, 124],\n", " [ 189, 256, 325],\n", " [ 396, 469, 544],\n", " [ 621, 700, 781]],\n", "\n", " [[ 864, 949, 1036],\n", " [1125, 1216, 1309],\n", " [1404, 1501, 1600],\n", " [1701, 1804, 1909]],\n", "\n", " [[2016, 2125, 2236],\n", " [2349, 2464, 2581],\n", " [2700, 2821, 2944],\n", " [3069, 3196, 3325]],\n", "\n", " [[3456, 3589, 3724],\n", " [3861, 4000, 4141],\n", " [4284, 4429, 4576],\n", " [4725, 4876, 5029]],\n", "\n", " [[5184, 5341, 5500],\n", " [5661, 5824, 5989],\n", " [6156, 6325, 6496],\n", " [6669, 6844, 7021]]], dtype=ht.int32, device=cpu:0, split=0)" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "b6abf056-e357819abb2ce1db79af7a26", "error": null, "execute_input": "other_dndarray = ht.arange(60,120, split=0).reshape(5,4,3) # distributed reshape\n\n# element-wise multiplication\ndndarray * other_dndarray\n", "execute_result": { "data": { "text/plain": "DNDarray([[[ 0, 61, 124],\n [ 189, 256, 325],\n [ 396, 469, 544],\n [ 621, 700, 781]],\n\n [[ 864, 949, 1036],\n [1125, 1216, 1309],\n [1404, 1501, 1600],\n [1701, 1804, 1909]],\n\n [[2016, 2125, 2236],\n [2349, 2464, 2581],\n [2700, 2821, 2944],\n [3069, 3196, 3325]],\n\n [[3456, 3589, 3724],\n [3861, 4000, 4141],\n [4284, 4429, 4576],\n [4725, 4876, 5029]],\n\n [[5184, 5341, 5500],\n [5661, 5824, 5989],\n [6156, 6325, 6496],\n [6669, 6844, 7021]]], dtype=ht.int32, device=cpu:0, split=0)" }, "execution_count": 32, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:53:09.351407Z" }, "output_type": "display_data" }, { "data": { "text/plain": [] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "44afd95b-fbaa3f38705a9e054856d875", "error": null, "execute_input": "other_dndarray = ht.arange(60,120, split=0).reshape(5,4,3) # distributed reshape\n\n# element-wise multiplication\ndndarray * other_dndarray\n", "execute_result": { "data": { "text/plain": "" }, "execution_count": 32, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:53:09.351908Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "other_dndarray = ht.arange(60,120, split=0).reshape(5,4,3) # distributed reshape\n", "\n", "# element-wise multiplication\n", "dndarray * other_dndarray\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "As we saw earlier, because the underlying data objects are PyTorch tensors, we can easily create DNDarrays on GPUs or move DNDarrays to GPUs. This allows us to perform distributed array operations on multi-GPU systems.\n", "\n", "So far we have demonstrated small, easy-to-parallelize arithmetical operations. Let's move to linear algebra. Heat's `linalg` module supports a wide range of linear algebra operations, including matrix multiplication. Matrix multiplication is a very common operation in data analysis, it is computationally intensive, and not trivial to parallelize. \n", "\n", "With Heat, you can perform matrix multiplication on distributed DNDarrays, and the operation will be parallelized across the MPI processes. Here on 4 GPUs:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%px\n", "# free up memory if necessary\n", "try:\n", " del x, y, z\n", "except NameError:\n", " pass\n", "\n", "n, m = 4000, 4000\n", "x = ht.random.randn(n, m, split=0, device=\"gpu\") # distributed RNG\n", "y = ht.random.randn(m, n, split=None, device=\"gpu\")\n", "z = x @ y\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "`ht.linalg.matmul` or `@` breaks down the matrix multiplication into a series of smaller `torch` matrix multiplications, which are then distributed across the MPI processes. This operation can be very communication-intensive on huge matrices that both require distribution, and users should choose the `split` axis carefully to minimize communication overhead." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can experiment with sizes and the `split` parameter (distribution axis) for both matrices and time the result. Note that:\n", "- If you set **`split=None` for both matrices**, each process (in this case, each GPU) will attempt to multiply the entire matrices. Depending on the matrix sizes, the GPU memory might be insufficient. (And if you can multiply the matrices on a single GPU, it's much more efficient to stick to PyTorch's `torch.linalg.matmul` function.)\n", "- If **`split` is not None for both matrices**, each process will only hold a slice of the data, and will need to communicate data with other processes in order to perform the multiplication. This **introduces huge communication overhead**, but allows you to perform the multiplication on larger matrices than would not fit in the memory of a single GPU.\n", "- If **`split` is None for one matrix and not None for the other**, the multiplication does not require communication, and the result will be distributed. If your data size allows it, you should always favor this option.\n", "\n", "Time the multiplication for different split parameters and see how the performance changes.\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[stdout:0] The slowest run took 46.16 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "1.63 ms ± 2.82 ms per loop (mean ± std. dev. of 5 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:2] The slowest run took 49.04 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "1.72 ms ± 3 ms per loop (mean ± std. dev. of 5 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:3] The slowest run took 43.62 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "1.55 ms ± 2.67 ms per loop (mean ± std. dev. of 5 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:1] The slowest run took 43.82 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "1.55 ms ± 2.67 ms per loop (mean ± std. dev. of 5 runs, 1 loop each)\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "%%px\n", "z = %timeit -n 1 -r 5 x @ y " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Heat supports many linear algebra operations:\n", "```bash\n", ">>> ht.linalg.\n", "ht.linalg.basics ht.linalg.hsvd_rtol( ht.linalg.projection( ht.linalg.triu(\n", "ht.linalg.cg( ht.linalg.inv( ht.linalg.qr( ht.linalg.vdot(\n", "ht.linalg.cross( ht.linalg.lanczos( ht.linalg.solver ht.linalg.vecdot(\n", "ht.linalg.det( ht.linalg.matmul( ht.linalg.svdtools ht.linalg.vector_norm(\n", "ht.linalg.dot( ht.linalg.matrix_norm( ht.linalg.trace( \n", "ht.linalg.hsvd( ht.linalg.norm( ht.linalg.transpose( \n", "ht.linalg.hsvd_rank( ht.linalg.outer( ht.linalg.tril( \n", "```\n", "\n", "and a lot more is in the works, including distributed eigendecompositions, SVD, and more. If the operation you need is not yet supported, leave us a note [here](https://github.com/helmholtz-analytics/heat/issues) and we'll get back to you." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can of course perform all operations on CPUs. You can leave out the `device` attribute entirely." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Interoperability\n", "\n", "We can easily create DNDarrays from PyTorch tensors and numpy ndarrays. We can also convert DNDarrays to PyTorch tensors and numpy ndarrays. This makes it easy to integrate Heat into existing PyTorch and numpy workflows. Here a basic example with xarrays:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "42eba82ca79146ccaa00e6abfc965d73", "version_major": 2, "version_minor": 0 }, "text/plain": [ "%px: 0%| | 0/4 [00:00\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
<xarray.DataArray (z: 1, y: 4, x: 3)> Size: 48B\n",
       "array([[[24, 25, 26],\n",
       "        [27, 28, 29],\n",
       "        [30, 31, 32],\n",
       "        [33, 34, 35]]], dtype=int32)\n",
       "Dimensions without coordinates: z, y, x
" ], "text/plain": [ "\u001b[0;31mOut[1:36]: \u001b[0m\n", " Size: 48B\n", "array([[[24, 25, 26],\n", " [27, 28, 29],\n", " [30, 31, 32],\n", " [33, 34, 35]]], dtype=int32)\n", "Dimensions without coordinates: z, y, x" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 1, "engine_uuid": "80de6265-3d236a5e1eb6ffa587e57739", "error": null, "execute_input": "import xarray as xr\n\nlocal_xr = xr.DataArray(dndarray.larray, dims=(\"z\", \"y\", \"x\"))\n# proceed with local xarray operations\nlocal_xr\n\n", "execute_result": { "data": { "text/html": "
\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
<xarray.DataArray (z: 1, y: 4, x: 3)> Size: 48B\narray([[[24, 25, 26],\n        [27, 28, 29],\n        [30, 31, 32],\n        [33, 34, 35]]], dtype=int32)\nDimensions without coordinates: z, y, x
", "text/plain": " Size: 48B\narray([[[24, 25, 26],\n [27, 28, 29],\n [30, 31, 32],\n [33, 34, 35]]], dtype=int32)\nDimensions without coordinates: z, y, x" }, "execution_count": 36, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:54:47.155658Z" }, "output_type": "display_data" }, { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
<xarray.DataArray (z: 2, y: 4, x: 3)> Size: 96B\n",
       "array([[[ 0,  1,  2],\n",
       "        [ 3,  4,  5],\n",
       "        [ 6,  7,  8],\n",
       "        [ 9, 10, 11]],\n",
       "\n",
       "       [[12, 13, 14],\n",
       "        [15, 16, 17],\n",
       "        [18, 19, 20],\n",
       "        [21, 22, 23]]], dtype=int32)\n",
       "Dimensions without coordinates: z, y, x
" ], "text/plain": [ "\u001b[0;31mOut[0:36]: \u001b[0m\n", " Size: 96B\n", "array([[[ 0, 1, 2],\n", " [ 3, 4, 5],\n", " [ 6, 7, 8],\n", " [ 9, 10, 11]],\n", "\n", " [[12, 13, 14],\n", " [15, 16, 17],\n", " [18, 19, 20],\n", " [21, 22, 23]]], dtype=int32)\n", "Dimensions without coordinates: z, y, x" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 0, "engine_uuid": "b6abf056-e357819abb2ce1db79af7a26", "error": null, "execute_input": "import xarray as xr\n\nlocal_xr = xr.DataArray(dndarray.larray, dims=(\"z\", \"y\", \"x\"))\n# proceed with local xarray operations\nlocal_xr\n\n", "execute_result": { "data": { "text/html": "
\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
<xarray.DataArray (z: 2, y: 4, x: 3)> Size: 96B\narray([[[ 0,  1,  2],\n        [ 3,  4,  5],\n        [ 6,  7,  8],\n        [ 9, 10, 11]],\n\n       [[12, 13, 14],\n        [15, 16, 17],\n        [18, 19, 20],\n        [21, 22, 23]]], dtype=int32)\nDimensions without coordinates: z, y, x
", "text/plain": " Size: 96B\narray([[[ 0, 1, 2],\n [ 3, 4, 5],\n [ 6, 7, 8],\n [ 9, 10, 11]],\n\n [[12, 13, 14],\n [15, 16, 17],\n [18, 19, 20],\n [21, 22, 23]]], dtype=int32)\nDimensions without coordinates: z, y, x" }, "execution_count": 36, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:54:47.155436Z" }, "output_type": "display_data" }, { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
<xarray.DataArray (z: 1, y: 4, x: 3)> Size: 48B\n",
       "array([[[48, 49, 50],\n",
       "        [51, 52, 53],\n",
       "        [54, 55, 56],\n",
       "        [57, 58, 59]]], dtype=int32)\n",
       "Dimensions without coordinates: z, y, x
" ], "text/plain": [ "\u001b[0;31mOut[3:36]: \u001b[0m\n", " Size: 48B\n", "array([[[48, 49, 50],\n", " [51, 52, 53],\n", " [54, 55, 56],\n", " [57, 58, 59]]], dtype=int32)\n", "Dimensions without coordinates: z, y, x" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 3, "engine_uuid": "44afd95b-fbaa3f38705a9e054856d875", "error": null, "execute_input": "import xarray as xr\n\nlocal_xr = xr.DataArray(dndarray.larray, dims=(\"z\", \"y\", \"x\"))\n# proceed with local xarray operations\nlocal_xr\n\n", "execute_result": { "data": { "text/html": "
\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
<xarray.DataArray (z: 1, y: 4, x: 3)> Size: 48B\narray([[[48, 49, 50],\n        [51, 52, 53],\n        [54, 55, 56],\n        [57, 58, 59]]], dtype=int32)\nDimensions without coordinates: z, y, x
", "text/plain": " Size: 48B\narray([[[48, 49, 50],\n [51, 52, 53],\n [54, 55, 56],\n [57, 58, 59]]], dtype=int32)\nDimensions without coordinates: z, y, x" }, "execution_count": 36, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:54:47.155903Z" }, "output_type": "display_data" }, { "data": { "text/html": [ "
\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "\n", "
<xarray.DataArray (z: 1, y: 4, x: 3)> Size: 48B\n",
       "array([[[36, 37, 38],\n",
       "        [39, 40, 41],\n",
       "        [42, 43, 44],\n",
       "        [45, 46, 47]]], dtype=int32)\n",
       "Dimensions without coordinates: z, y, x
" ], "text/plain": [ "\u001b[0;31mOut[2:36]: \u001b[0m\n", " Size: 48B\n", "array([[[36, 37, 38],\n", " [39, 40, 41],\n", " [42, 43, 44],\n", " [45, 46, 47]]], dtype=int32)\n", "Dimensions without coordinates: z, y, x" ] }, "metadata": { "after": null, "completed": null, "data": {}, "engine_id": 2, "engine_uuid": "ab9782c4-a0794501cf2af7341aeff65d", "error": null, "execute_input": "import xarray as xr\n\nlocal_xr = xr.DataArray(dndarray.larray, dims=(\"z\", \"y\", \"x\"))\n# proceed with local xarray operations\nlocal_xr\n\n", "execute_result": { "data": { "text/html": "
\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
<xarray.DataArray (z: 1, y: 4, x: 3)> Size: 48B\narray([[[36, 37, 38],\n        [39, 40, 41],\n        [42, 43, 44],\n        [45, 46, 47]]], dtype=int32)\nDimensions without coordinates: z, y, x
", "text/plain": " Size: 48B\narray([[[36, 37, 38],\n [39, 40, 41],\n [42, 43, 44],\n [45, 46, 47]]], dtype=int32)\nDimensions without coordinates: z, y, x" }, "execution_count": 36, "metadata": {} }, "follow": null, "msg_id": null, "outputs": [], "received": null, "started": null, "status": null, "stderr": "", "stdout": "", "submitted": "2025-11-21T13:54:47.155715Z" }, "output_type": "display_data" } ], "source": [ "%%px\n", "import xarray as xr\n", "\n", "local_xr = xr.DataArray(dndarray.larray, dims=(\"z\", \"y\", \"x\"))\n", "# proceed with local xarray operations\n", "local_xr\n", "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**NOTE:** this is not a distributed `xarray`, but local xarray objects on each rank.\n", "Work on [expanding xarray support](https://github.com/helmholtz-analytics/heat/pull/1183) is ongoing.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Heat will try to reuse the memory of the original array as much as possible. If you would prefer a copy with different memory, the ```copy``` keyword argument can be used when creating a DNDArray from other libraries." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[stdout:2] tensor([-1, 1, 2, 3, 4])\n", "tensor([0, 1, 2, 3, 4])\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:3] tensor([-1, 1, 2, 3, 4])\n", "tensor([0, 1, 2, 3, 4])\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:1] tensor([-1, 1, 2, 3, 4])\n", "tensor([0, 1, 2, 3, 4])\n" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "[stdout:0] tensor([-1, 1, 2, 3, 4])\n", "tensor([0, 1, 2, 3, 4])\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "%%px\n", "import torch\n", "torch_array = torch.arange(5)\n", "heat_array = ht.array(torch_array, copy=False)\n", "heat_array[0] = -1\n", "print(torch_array)\n", "\n", "torch_array = torch.arange(5)\n", "heat_array = ht.array(torch_array, copy=True)\n", "heat_array[0] = -1\n", "print(torch_array)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Interoperability is a key feature of Heat, and we are constantly working to increase Heat's compliance with the [Python array API standard](https://data-apis.org/array-api/latest/). As usual, please [let us know](https://github.com/helmholtz-analytics/heat/issues) if you encounter any issues or have any feature requests." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Exercise: Array Manipulation (Basic and Distributed)\n", "\n", "This exercise shows the difference between a standard array and a distributed array when performing operations.\n", "\n", "**Part 1: Basic Operations (split=None, sample code below)**\n", "\n", "Task: Create a 10×5 DNDarray of random integers between 0 and 9 (inclusive).\n", "\n", "- Calculate its transpose (A.T​).\n", "\n", "- Find the maximum value along the second axis (axis=1).\n", "\n", "- If you have a GPU available, try creating the same array on the GPU and repeat the operations.\n", "\n", "**Part 2: Distributed Operations (split=0)**\n", "\n", "Task: Repeat the steps, but distribute the original array A along its first axis (split=0).\n", "\n", "**NB:** Notice how the `split` property now changes after transpose (from 0 to 1) and is removed (None) after the reduction over the split axis (axis=1). The lshape also shows the local data chunk size." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[stdout:0] --- Part 1: Basic Operations (split=None) ---\n", "A: Shape=(10, 5), Split=None, Device=cpu:0\n", "A_T: Shape=(5, 10), Split=None, Device=cpu:0\n", "Max Vals: Shape=(10,), Split=None, Device=cpu:0\n", "Max Vals Result (A's Maxima):\n", "DNDarray([7, 8, 5, 8, 5, 9, 7, 9, 9, 7], dtype=ht.int32, device=cpu:0, split=None)\n" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "%%px\n", "### Exercise Part 1, sample code\n", "\n", "import heat as ht\n", "\n", "# Task 1: Create a 10x5 DNDarray (no split/split=None)\n", "# NOTE: Change 'cpu' to 'gpu' if you have one!\n", "A = ht.random.randint(0, 10, size=(10, 5))\n", "\n", "# Task 2: Calculate the transpose 'A_T'\n", "A_T = A.T\n", "\n", "# Task 3: Find the maximum value along axis=1\n", "max_vals = ht.max(A, axis=1)\n", "\n", "# Check: Print properties and results (only on rank 0)\n", "if A.comm.rank == 0:\n", " print(\"--- Part 1: Basic Operations (split=None) ---\")\n", " print(f\"A: Shape={A.shape}, Split={A.split}, Device={A.device}\")\n", " print(f\"A_T: Shape={A_T.shape}, Split={A_T.split}, Device={A_T.device}\")\n", " print(f\"Max Vals: Shape={max_vals.shape}, Split={max_vals.split}, Device={max_vals.device}\")\n", " print(f\"Max Vals Result (A's Maxima):\\n{max_vals}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%px\n", "## Exercise Part 2, your code here" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "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.10.12" } }, "nbformat": 4, "nbformat_minor": 4 }