Repository: e3nn/e3nn (language: python)

Issue:
🌟 Show short names in the docs.
@mariogeiger asked in Slack (edited for clarity):
> Do you know how to change `e3nn.o3.irreps.Irrep` to `e3nn.o3.Irrep` in the docs? Because `irreps.py` is the name of the file but I want the user to import from `o3` since `__init__.py` of `o3` imports the contents of `irreps.py`.

I'm creating this issue and assigning myself to look into it. (Edit: Oops, I don't have privileges to assign myself to this issue. But someone else can assign me.)


Discussion / hints:

Fix patch:
diff --git a/INSTALL.md b/INSTALL.md
index ac7cf909..7ec78964 100644
--- a/INSTALL.md
+++ b/INSTALL.md
@@ -25,7 +25,7 @@ See [here](https://github.com/rusty1s/pytorch_geometric#installation) to get cud
 
 ### Optional: opt_einsum_fx (beta)
 
-e3nn can use [`opt_einsum_fx`]() to optimize the performance of `TensorProduct`s. To enable this, install `opt_einsum_fx`:
+e3nn can use [`opt_einsum_fx`](https://github.com/Linux-cpp-lisp/opt_einsum_fx) to optimize the performance of `TensorProduct`s. To enable this, install `opt_einsum_fx`:
 ```bash
 $ git clone https://github.com/Linux-cpp-lisp/opt_einsum_fx.git
 $ cd opt_einsum_fx/
diff --git a/docs/examples/tetris_polynomial.rst b/docs/examples/tetris_polynomial.rst
index 3dea01b2..f70034a4 100644
--- a/docs/examples/tetris_polynomial.rst
+++ b/docs/examples/tetris_polynomial.rst
@@ -19,13 +19,13 @@ And the following features of `pytorch_geometric <https://pytorch-geometric.read
 .. rubric:: the model
 
 .. literalinclude:: ../../examples/tetris_polynomial.py
-    :lines: 57-94
+    :lines: 57-103
 
 
 .. rubric:: training
 
 .. literalinclude:: ../../examples/tetris_polynomial.py
-    :lines: 98-118
+    :lines: 107-127
     :dedent: 4
 
 Full code `here <https://github.com/e3nn/e3nn/blob/main/examples/tetris_polynomial.py>`_
diff --git a/docs/guide/jit.rst b/docs/guide/jit.rst
index c16ebe3e..4f7997a3 100644
--- a/docs/guide/jit.rst
+++ b/docs/guide/jit.rst
@@ -2,7 +2,7 @@
 TorchScript JIT Support
 =======================
 
-PyTorch provides two ways to compile code into TorchScript: ``tracing and scripting <https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html>``_. Tracing follows the tensor operations on an example input, allowing complex Python control flow if that control flow does not depend on the data itself. Scripting compiles a subset of Python directly into TorchScript, allowing data-dependent control flow but only limited Python features.
+PyTorch provides two ways to compile code into TorchScript: `tracing and scripting <https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html>`_. Tracing follows the tensor operations on an example input, allowing complex Python control flow if that control flow does not depend on the data itself. Scripting compiles a subset of Python directly into TorchScript, allowing data-dependent control flow but only limited Python features.
 
 This is a problem for e3nn, where many modules --- such as ``TensorProduct`` --- use significant Python control flow based on ``Irreps`` as well as features like inheritance that are incompatible with scripting. Other modules like ``Gate``, however, contain important but simple data-dependent control flow. Thus ``Gate`` needs to be scripted, even though it contains a ``TensorProduct`` that has to be traced.
 
@@ -28,7 +28,7 @@ We define a simple module that includes data-dependent control flow:
             if torch.any(norm > 7.):
                 return norm
             else:
-                return norm * 0.5 
+                return norm * 0.5
 
     irreps = Irreps("2x0e + 1x1o")
     mod = MyModule(irreps)
@@ -88,7 +88,7 @@ Say we define:
         def forward(self, x):
             return self.mymod(x) + 3.
 
-And trace an instance of `AnotherModule` using `e3nn.util.jit.trace`:
+And trace an instance of ``AnotherModule`` using `e3nn.util.jit.trace`:
 
 .. jupyter-execute::
 
@@ -117,7 +117,7 @@ Customizing Tracing Inputs
 ==========================
 
 Submodules can also be compiled automatically using tracing if they are marked with ``@compile_mode('trace')``. When submodules are compiled by tracing it must be possible to generate plausible input examples on the fly.
-    
+
 These example inputs can be generated automatically based on the ``irreps_in`` of the module (the specifics are the same as for ``assert_equivariant``). If this is not possible or would yield incorrect results, a module can define a ``_make_tracing_inputs`` method that generates example inputs of correct shape and type.
 
 .. jupyter-execute::
@@ -127,18 +127,18 @@ These example inputs can be generated automatically based on the ``irreps_in`` o
         def forward(self, x: torch.Tensor, indexes: torch.LongTensor):
             return x[indexes].sum()
 
-        # Because this module has no `irreps_in`, and because 
-        # `irreps_in` can't describe indexes, since it's a LongTensor, 
+        # Because this module has no `irreps_in`, and because
+        # `irreps_in` can't describe indexes, since it's a LongTensor,
         # we impliment _make_tracing_inputs
         def _make_tracing_inputs(self, n: int):
             import random
-            # The compiler asks for n example inputs --- 
-            # this is only a suggestion, the only requirement 
+            # The compiler asks for n example inputs ---
+            # this is only a suggestion, the only requirement
             # is that at least one be returned.
             return [
                 {
                     'forward': (
-                        torch.randn(5, random.randint(1, 3)), 
+                        torch.randn(5, random.randint(1, 3)),
                         torch.arange(3)
                     )
                 }


Test patch:
diff --git a/docs/guide/equivar_testing.rst b/docs/guide/equivar_testing.rst
index 7b36df44..350be593 100644
--- a/docs/guide/equivar_testing.rst
+++ b/docs/guide/equivar_testing.rst
@@ -10,12 +10,12 @@ In `e3nn.util.test`, the library provides some tools for confirming that functio
     In [3]: from e3nn.util.test import equivariance_error
 
     In [4]: equivariance_error(
-                tp, 
+                tp,
                 args_in=[tp.irreps_in1.randn(1, -1), tp.irreps_in2.randn(1, -1)],
                 irreps_in=[tp.irreps_in1, tp.irreps_in2],
                 irreps_out=[tp.irreps_out]
             )
-    Out[4]: 
+    Out[4]:
     {(tensor(0.), False): tensor(1.1921e-07, grad_fn=<MaxBackward1>),
     (tensor(1.), False): tensor(1.1921e-07, grad_fn=<MaxBackward1>)}
 
@@ -53,7 +53,7 @@ These can be used to test models that operate on full graphs that include positi
         irreps_out=[f.irreps_out],
     )
 
-To test equivariance on a specific graph, `args_in` can be used::
+To test equivariance on a specific graph, ``args_in`` can be used::
 
     assert_equivariant(
         wrapper,