miscellaneous fixes

Author: wjakob

docs/nn.rst

@@ -3,7 +3,7 @@
 .. _neural_nets:
 
 Neural Networks
----------------
+===============
 
 Dr.Jit's neural network infrastructure builds on :ref:`cooperative vectors
 <coop_vec>`. Please review their documentation before reading this section.
@@ -26,16 +26,16 @@ List
 
 The set of neural network module currently includes:
 
-- Sequential evaluation of a list of models: :py:class:`nn.Sequential`.
+- Sequential evaluation of a list of models: :py:class:`nn.Sequential <Sequential>`.
 
-- Linear/affine layers: :py:class:`nn.Linear`.
+- Linear/affine layers: :py:class:`nn.Linear <Linear>`.
 
-- Encoding layers: :py:class:`nn.SinEncode`, :py:class:`nn.TriEncode`.
+- Encoding layers: :py:class:`nn.SinEncode <SinEncode>`, :py:class:`nn.TriEncode <TriEncode>`.
 
-- Activation functions and other nonlinear transformations: :py:class:`ReLU`, :py:class:`LeakyReLU`,
-  :py:class:`Exp`, :py:class:`Exp2`, :py:class:`Tanh`.
+- Activation functions and other nonlinear transformations: :py:class:`nn.ReLU <ReLU>`, :py:class:`nn.LeakyReLU <LeakyReLU>`,
+  :py:class:`nn.Exp <nn.Exp>`, :py:class:`nn.Exp2 <Exp2>`, :py:class:`nn.Tanh <Tanh>`.
 
-- Miscellaneous: :py:class:`nn.Cast`, :py:class:`nn.ScaleAdd`.
+- Miscellaneous: :py:class:`nn.Cast <Cast>`, :py:class:`nn.ScaleAdd <ScaleAdd>`.
 
 Example
 -------
@@ -48,9 +48,13 @@ Great Wave off Kanagawa
 <https://en.wikipedia.org/wiki/The_Great_Wave_off_Kanagawa>`__ (left).
 
 .. image:: https://rgl.s3.eu-central-1.amazonaws.com/media/uploads/wjakob/2024/06/coopvec-screenshot.png
-  :width: 300
+  :width: 600
   :align: center
 
+The optimization uses the :py:class:`dr.opt.Adam <drjit.opt.Adam>` optimizer
+and :py:class:`dr.opt.GradScaler <drjit.opt.GradScaler>` gradient scaler for
+adaptive mixed-precision training.
+
 .. code-block:: python
 
     from tqdm.auto import tqdm
@@ -126,3 +130,4 @@ Great Wave off Kanagawa
     ax[1].imshow(dr.clip(img, 0, 1))
     fig.tight_layout()
     plt.show()
+

drjit/nn.py

@@ -1,6 +1,11 @@
 from __future__ import annotations
 import drjit
-from typing import Tuple, Sequence, Union, Type, TypeAlias, Optional, Any
+import sys
+
+if sys.version_info < (3, 11):
+    from typing_extensions import Tuple, Sequence, Union, Type, TypeAlias, Optional, Any
+else:
+    from typing import Tuple, Sequence, Union, Type, TypeAlias, Optional, Any
 
 # Import classes/functions from C++ extension
 MatrixView = drjit.detail.nn.MatrixView

src/python/coop_vec.cpp

@@ -447,13 +447,13 @@ static nb::object repack_impl(const char *name, MatrixLayout layout,
         nb::list result;
         for (nb::handle h : l)
             result.append(repack_impl(name, layout, h, offset, items));
-        return result;
+        return std::move(result);
     } else if (arg_tp.is(&PyDict_Type)) {
         nb::dict d = nb::borrow<nb::dict>(arg);
         nb::dict result;
         for (auto [k, v] : d)
             result[k] = repack_impl(name, layout, v, offset, items);
-        return result;
+        return std::move(result);
     } else if (nb::dict ds = get_drjit_struct(arg_tp); ds.is_valid()) {
         nb::object tmp = arg_tp();
         for (auto [k, v] : ds)