.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "tutorials/tensordict_preallocation.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_tutorials_tensordict_preallocation.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_tutorials_tensordict_preallocation.py:


Pre-allocating memory with TensorDict
=====================================
**Author**: `Tom Begley <https://github.com/tcbegley>`_

In this tutorial you will learn how to take advantage of memory pre-allocation in
:class:`~.TensorDict`.

.. GENERATED FROM PYTHON SOURCE LINES 11-12

Suppose that we have a function that returns a :class:`~.TensorDict`

.. GENERATED FROM PYTHON SOURCE LINES 12-22

.. code-block:: Python



    import torch
    from tensordict.tensordict import TensorDict


    def make_tensordict():
        return TensorDict({"a": torch.rand(3), "b": torch.rand(3, 4)}, [3])









.. GENERATED FROM PYTHON SOURCE LINES 28-30

Perhaps we want to call this function multiple times and use the results to populate
a single :class:`~.TensorDict`.

.. GENERATED FROM PYTHON SOURCE LINES 30-39

.. code-block:: Python


    N = 10
    tensordict = TensorDict({}, batch_size=[N, 3])

    for i in range(N):
        tensordict[i] = make_tensordict()

    print(tensordict)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    TensorDict(
        fields={
            a: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
            b: Tensor(shape=torch.Size([10, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
        batch_size=torch.Size([10, 3]),
        device=None,
        is_shared=False)




.. GENERATED FROM PYTHON SOURCE LINES 40-51

Because we have specified the ``batch_size`` of ``tensordict``, during the first
iteration of the loop we populate ``tensordict`` with empty tensors whose first
dimension is size ``N``, and whose remaining dimensions are determined by the return
value of ``make_tensordict``. In the above example, we pre-allocate an array of zeros
of size ``torch.Size([10, 3])`` for the key ``"a"``, and an array size
``torch.Size([10, 3, 4])`` for the key ``"b"``. Subsequent iterations of the loop are
written in place. As a result, if not all values are filled, they get the default
value of zero.

Let us demonstrate what is going on by stepping through the above loop. We first
initialise an empty :class:`~.TensorDict`.

.. GENERATED FROM PYTHON SOURCE LINES 51-56

.. code-block:: Python


    N = 10
    tensordict = TensorDict({}, batch_size=[N, 3])
    print(tensordict)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    TensorDict(
        fields={
        },
        batch_size=torch.Size([10, 3]),
        device=None,
        is_shared=False)




.. GENERATED FROM PYTHON SOURCE LINES 57-60

After the first iteration, ``tensordict`` has been prepopulated with tensors for both
``"a"`` and ``"b"``. These tensors contain zeros except for the first row which we
have assigned random values to.

.. GENERATED FROM PYTHON SOURCE LINES 60-69

.. code-block:: Python


    random_tensordict = make_tensordict()
    tensordict[0] = random_tensordict

    assert (tensordict[1:] == 0).all()
    assert (tensordict[0] == random_tensordict).all()

    print(tensordict)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    TensorDict(
        fields={
            a: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
            b: Tensor(shape=torch.Size([10, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
        batch_size=torch.Size([10, 3]),
        device=None,
        is_shared=False)




.. GENERATED FROM PYTHON SOURCE LINES 70-71

Subsequent iterations, we update the pre-allocated tensors in-place.

.. GENERATED FROM PYTHON SOURCE LINES 71-79

.. code-block:: Python


    a = tensordict["a"]
    random_tensordict = make_tensordict()
    tensordict[1] = random_tensordict

    # the same tensor is stored under "a", but the values have been updated
    assert tensordict["a"] is a
    assert (tensordict[:2] != 0).all()








.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 0.003 seconds)


.. _sphx_glr_download_tutorials_tensordict_preallocation.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: tensordict_preallocation.ipynb <tensordict_preallocation.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: tensordict_preallocation.py <tensordict_preallocation.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: tensordict_preallocation.zip <tensordict_preallocation.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_