diff --git a/docs/tutorials/etp.rst b/docs/tutorials/etp.rst index 8efd6fd..08c15ef 100644 --- a/docs/tutorials/etp.rst +++ b/docs/tutorials/etp.rst @@ -16,7 +16,7 @@ Equivariant Tensor Product ========================== -The submodule :class:`cuequivariance.descriptors` contains many descriptors of Equivariant Tensor Products (:class:`cuequivariance.EquivariantTensorProduct`). +The submodule :class:`cuequivariance.descriptors` contains many descriptors of Equivariant Tensor Products represented by the class :class:`cuequivariance.EquivariantTensorProduct`. Examples -------- @@ -59,6 +59,7 @@ Execution on JAX import jax import jax.numpy as jnp + import cuequivariance as cue import cuequivariance_jax as cuex e = cue.descriptors.linear( @@ -83,6 +84,7 @@ We can execute an :class:`cuequivariance.EquivariantTensorProduct` with PyTorch. .. jupyter-execute:: import torch + import cuequivariance as cue import cuequivariance_torch as cuet e = cue.descriptors.linear( diff --git a/docs/tutorials/stp.rst b/docs/tutorials/stp.rst index b6bfbb0..95ed850 100644 --- a/docs/tutorials/stp.rst +++ b/docs/tutorials/stp.rst @@ -50,7 +50,7 @@ The subscripts of this tensor product are "uv,iu,iv" where "uv" represents the m d Each operand of the tensor product descriptor has a list of segments. -We can add segments to the descriptor using the `add_segment` method. +We can add segments to the descriptor using the :meth:`add_segment <cuequivariance.SegmentedTensorProduct.add_segment>` method. We can add the segments of the input and output representations to the descriptor. .. jupyter-execute:: @@ -90,7 +90,7 @@ Finally, we can normalize the paths for the last operand such that the output is As we can see, the paths coefficients has been normalized. -Now we can create a tensor product from the descriptor and execute it. In PyTorch, we can use the :class:`cuet.TensorProduct` class. +Now we can create a tensor product from the descriptor and execute it. In PyTorch, we can use the :class:`cuet.TensorProduct <cuequivariance_torch.TensorProduct>` class. .. jupyter-execute:: @@ -98,7 +98,7 @@ Now we can create a tensor product from the descriptor and execute it. In PyTorc linear_torch -In JAX, we can use the :func:`cuex.tensor_product` function. +In JAX, we can use the :func:`cuex.tensor_product <cuequivariance_jax.tensor_product>` function. .. jupyter-execute::