[BugFix, Doc] Fix tutorial (#606)

.github/workflows/docs.yml

@@ -4,6 +4,9 @@ on:
   push:
     branches:
       - main
+  pull_request:
+    branches:
+      - "*"
   workflow_dispatch:
 
 concurrency:
@@ -14,46 +17,126 @@ concurrency:
 
 jobs:
   build-docs:
-    runs-on: ubuntu-20.04
-    defaults:
-      run:
-        shell: bash -l {0}
-    steps:
-      - name: Checkout tensordict
-        uses: actions/checkout@v3
-      - name: Setup Miniconda
-        uses: conda-incubator/setup-miniconda@v2
-        with:
-          activate-environment: build_docs
-          python-version: "3.9"
-      - name: Check Python version
-        run: |
-          python --version
-      - name: Install PyTorch
-        run: |
-          pip3 install pip --upgrade
-          pip3 install packaging
-          pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
-      - name: Install TensorDict
-        run: |
-          python -m pip install -e .
-      - name: Test tensordict installation
-        run: |
-          mkdir _tmp
-          cd _tmp
-          python -c "import tensordict"
-          cd ..
-      - name: Build the docset
-        working-directory: ./docs
-        run: |
-          python -m pip install -r requirements.txt
-          make docs
-      - name: Get output time
-        run: echo "The time was ${{ steps.build.outputs.time }}"
-      - name: Deploy
-        uses: JamesIves/github-pages-deploy-action@releases/v3
-        with:
-          ACCESS_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          BRANCH: gh-pages # The branch the action should deploy to.
-          FOLDER: docs/build/html # The folder the action should deploy.
-          CLEAN: false
+    strategy:
+      matrix:
+        python_version: ["3.9"]
+        cuda_arch_version: ["12.1"]
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      repository: pytorch/tensordict
+      upload-artifact: docs
+      runner: "linux.g5.4xlarge.nvidia.gpu"
+      docker-image: "nvidia/cudagl:11.4.0-base"
+      timeout: 120
+      script: |
+        set -e
+        set -v
+        apt-get update && apt-get install -y git wget gcc g++
+        root_dir="$(pwd)"
+        conda_dir="${root_dir}/conda"
+        env_dir="${root_dir}/env"
+        os=Linux
+        
+        # 1. Install conda at ./conda
+        printf "* Installing conda\n"
+        wget -O miniconda.sh "http://repo.continuum.io/miniconda/Miniconda3-latest-${os}-x86_64.sh"
+        bash ./miniconda.sh -b -f -p "${conda_dir}"
+        eval "$(${conda_dir}/bin/conda shell.bash hook)"
+        printf "* Creating a test environment\n"
+        conda create --prefix "${env_dir}" -y python=3.8
+        printf "* Activating\n"
+        conda activate "${env_dir}"
+        
+        # 2. upgrade pip, ninja and packaging
+        apt-get install python3.8 python3-pip unzip -y
+        python3 -m pip install --upgrade pip
+        python3 -m pip install setuptools ninja packaging -U
+        
+        # 3. check python version
+        python3 --version
+        
+        # 4. Check git version
+        git version
+        
+        # 5. Install PyTorch
+        python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu --quiet --root-user-action=ignore
+        
+        # 6. Install tensordict
+        python3 setup.py develop
+        
+        # 7. Install requirements
+        python3 -m pip install -r docs/requirements.txt --quiet --root-user-action=ignore
+        
+        # 8. Test tensordict installation
+        mkdir _tmp
+        cd _tmp
+        PYOPENGL_PLATFORM=egl MUJOCO_GL=egl python3 -c """from tensordict import *"""
+        cd ..
+        
+        # 10. Build doc
+        cd ./docs
+        make docs
+        # PYOPENGL_PLATFORM=egl MUJOCO_GL=egl sphinx-build ./source _local_build
+        cd ..
+
+        cp -r docs/build/* "${RUNNER_ARTIFACT_DIR}"
+        echo $(ls "${RUNNER_ARTIFACT_DIR}")
+        if [[ ${{ github.event_name == 'pull_request' }} ]]; then
+          cp -r docs/build/* "${RUNNER_DOCS_DIR}"
+        fi
+
+  upload:
+    needs: build-docs
+    if: github.repository == 'pytorch/tensordict' && github.event_name == 'push' && 
+        ((github.ref_type == 'branch' && github.ref_name == 'main') || github.ref_type == 'tag')
+    permissions:
+      contents: write
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      repository: pytorch/tensordict
+      download-artifact: docs
+      ref: gh-pages
+      test-infra-ref: main
+      script: |
+        set -euo pipefail
+
+        REF_TYPE=${{ github.ref_type }}
+        REF_NAME=${{ github.ref_name }}
+
+        # TODO: adopt this behaviour
+        # if [[ "${REF_TYPE}" == branch ]]; then
+        #   TARGET_FOLDER="${REF_NAME}"
+        # elif [[ "${REF_TYPE}" == tag ]]; then
+        #   case "${REF_NAME}" in
+        #     *-rc*)
+        #       echo "Aborting upload since this is an RC tag: ${REF_NAME}"
+        #       exit 0
+        #       ;;
+        #     *)
+        #       # Strip the leading "v" as well as the trailing patch version. For example:
+        #       # 'v0.15.2' -> '0.15'
+        #       TARGET_FOLDER=$(echo "${REF_NAME}" | sed 's/v\([0-9]\+\)\.\([0-9]\+\)\.[0-9]\+/\1.\2/')
+        #       ;;
+        #   esac
+        # fi
+        TARGET_FOLDER="./"
+        echo "Target Folder: ${TARGET_FOLDER}"
+
+        # mkdir -p "${TARGET_FOLDER}"
+        # rm -rf "${TARGET_FOLDER}"/*
+        echo $(ls "${RUNNER_ARTIFACT_DIR}")
+        rsync -a "${RUNNER_ARTIFACT_DIR}"/ "${TARGET_FOLDER}"
+        git add "${TARGET_FOLDER}" || true
+
+        # if [[ "${TARGET_FOLDER}" == main ]]; then
+        #   mkdir -p _static
+        #   rm -rf _static/*
+        #   cp -r "${TARGET_FOLDER}"/_static/* _static
+        #   git add _static || true
+        # fi
+
+        git config user.name 'pytorchbot'
+        git config user.email '[email protected]'
+        git config http.postBuffer 524288000
+        git commit -m "auto-generating sphinx docs" || true
+        git push

docs/source/reference/nn.rst

@@ -320,7 +320,7 @@ first traced using :func:`~.symbolic_trace`.
 Distributions
 -------------
 
-.. py:currentmodule::tensordict.nn.distributions
+.. currentmodule:: tensordict.nn.distributions
 
 .. autosummary::
     :toctree: generated/

docs/source/reference/tensordict.rst

@@ -50,7 +50,7 @@ Memory-mapped tensors
 `tensordict` offers the :class:`~tensordict.MemoryMappedTensor` primitive which
 allows you to work with tensors stored in physical memory in a handy way.
 The main advantages of :class:`~tensordict.MemoryMappedTensor`
-are its easiness of construction (no need to handle the storage of a tensor),
+are its ease of construction (no need to handle the storage of a tensor),
 the possibility to work with big contiguous data that would not fit in memory,
 an efficient (de)serialization across processes and efficient indexing of
 stored tensors.
@@ -96,3 +96,4 @@ Utils
     pad_sequence
     dense_stack_tds
     set_lazy_legacy
+    lazy_legacy

tensordict/__init__.py

@@ -17,17 +17,22 @@
 from tensordict.memmap_deprec import is_memmap, MemmapTensor, set_transfer_ownership
 from tensordict.persistent import PersistentTensorDict
 from tensordict.tensorclass import NonTensorData, tensorclass
-from tensordict.utils import assert_allclose_td, is_batchedtensor, is_tensorclass
+from tensordict.utils import (
+    assert_allclose_td,
+    is_batchedtensor,
+    is_tensorclass,
+    lazy_legacy,
+    set_lazy_legacy,
+)
+from tensordict._pytree import *
+
+from tensordict._tensordict import unravel_key, unravel_key_list
 
 try:
     from tensordict.version import __version__
 except ImportError:
     __version__ = None
 
-from tensordict._pytree import *
-
-from tensordict._tensordict import unravel_key, unravel_key_list
-
 __all__ = [
     "LazyStackedTensorDict",
     "MemmapTensor",

tensordict/_td.py

@@ -554,7 +554,7 @@ def __setitem__(
                 if key in keys:
                     self._set_at_str(key, item, index, validated=False)
                 else:
-                    subtd.set(key, item)
+                    subtd.set(key, item, inplace=True)
         else:
             for key in self.keys():
                 self.set_at_(key, value, index)

tutorials/sphinx_tuto/data_fashion.py

@@ -16,7 +16,7 @@
 import torch
 import torch.nn as nn
 
-from tensordict import MemmapTensor, TensorDict
+from tensordict import MemoryMappedTensor, TensorDict
 from torch.utils.data import DataLoader
 from torchvision import datasets
 from torchvision.transforms import ToTensor
@@ -50,27 +50,26 @@
 # batches of transformed data from disk rather than repeatedly load and transform
 # individual images.
 #
-# First we create the ``MemmapTensor`` containers.
+# First we create the :class:`~tensordict.MemoryMappedTensor` containers.
 
 
 training_data_td = TensorDict(
     {
-        "images": MemmapTensor(
-            len(training_data),
-            *training_data[0][0].squeeze().shape,
+        "images": MemoryMappedTensor.empty(
+            (len(training_data), *training_data[0][0].squeeze().shape),
             dtype=torch.float32,
         ),
-        "targets": MemmapTensor(len(training_data), dtype=torch.int64),
+        "targets": MemoryMappedTensor.empty((len(training_data),), dtype=torch.int64),
     },
     batch_size=[len(training_data)],
     device=device,
 )
 test_data_td = TensorDict(
     {
-        "images": MemmapTensor(
-            len(test_data), *test_data[0][0].squeeze().shape, dtype=torch.float32
+        "images": MemoryMappedTensor.empty(
+            (len(test_data), *test_data[0][0].squeeze().shape), dtype=torch.float32
         ),
-        "targets": MemmapTensor(len(test_data), dtype=torch.int64),
+        "targets": MemoryMappedTensor.empty((len(test_data),), dtype=torch.int64),
     },
     batch_size=[len(test_data)],
     device=device,

tutorials/sphinx_tuto/tensorclass_fashion.py

@@ -17,7 +17,7 @@
 import torch
 import torch.nn as nn
 
-from tensordict import MemmapTensor
+from tensordict import MemoryMappedTensor
 from tensordict.prototype import tensorclass
 from torch.utils.data import DataLoader
 from torchvision import datasets
@@ -68,10 +68,10 @@ class FashionMNISTData:
     @classmethod
     def from_dataset(cls, dataset, device=None):
         data = cls(
-            images=MemmapTensor(
-                len(dataset), *dataset[0][0].squeeze().shape, dtype=torch.float32
+            images=MemoryMappedTensor.empty(
+                (len(dataset), *dataset[0][0].squeeze().shape), dtype=torch.float32
             ),
-            targets=MemmapTensor(len(dataset), dtype=torch.int64),
+            targets=MemoryMappedTensor.empty((len(dataset),), dtype=torch.int64),
             batch_size=[len(dataset)],
             device=device,
         )