Updated documentation.

Author: vbharadwaj-bk

docs/api.rst

@@ -8,7 +8,7 @@ OpenEquivariance API
 OpenEquivariance exposes two key classes: :py:class:`openequivariance.TensorProduct`, which replaces
 ``o3.TensorProduct`` from e3nn, and :py:class:`openequivariance.TensorProductConv`, which fuses
 the CG tensor product with a subsequent graph convolution. Initializing either class triggers
-JIT compilation of a custom kernel, which can take a few seconds. 
+JIT compilation of a custom kernel, which can take a few seconds.
 
 Both classes require a configuration object specified 
 by :py:class:`openequivariance.TPProblem`, which has a constructor
@@ -17,6 +17,9 @@ We recommend reading the `e3nn documentation <https://docs.e3nn.org/en/latest/>`
 trying our code. OpenEquivariance cannot accelerate all tensor products; see 
 :doc:`this page </supported_ops>` for a list of supported configurations.
 
+PyTorch API
+------------------------
+
 .. autoclass:: openequivariance.TensorProduct
     :members: forward, reorder_weights_from_e3nn, reorder_weights_to_e3nn, to
     :undoc-members:
@@ -27,19 +30,21 @@ trying our code. OpenEquivariance cannot accelerate all tensor products; see
     :undoc-members:
     :exclude-members: name
 
-.. autoclass:: openequivariance.TPProblem
-    :members:
-    :undoc-members:
-
 .. autofunction:: openequivariance.torch_to_oeq_dtype
 
 .. autofunction:: openequivariance.torch_ext_so_path
 
-OpenEquivariance JAX API
+JAX API
 ------------------------
 The JAX API consists of ``TensorProduct`` and ``TensorProductConv``
 classes that behave identically to their PyTorch counterparts. These classes
-do not conform exactly to the e3nn-jax API, but perform the same computation. 
+do not conform exactly to the e3nn-jax API, but perform the same computation.
+
+If you plan to use ``oeq.jax`` without PyTorch installed, 
+you need to set ``OEQ_NOTORCH=1`` in your local environment (within Python,
+``os.environ["OEQ_NOTORCH"] = 1``). For the moment, we require this to avoid 
+breaking the PyTorch version of OpenEquivariance.
+
 
 .. autoclass:: openequivariance.jax.TensorProduct
     :members: forward, reorder_weights_from_e3nn, reorder_weights_to_e3nn
@@ -51,6 +56,12 @@ do not conform exactly to the e3nn-jax API, but perform the same computation.
     :undoc-members:
     :exclude-members: 
 
+Common API
+---------------------
+
+.. autoclass:: openequivariance.TPProblem
+    :members:
+    :undoc-members:
 
 API Identical to e3nn
 ---------------------

docs/installation.rst

@@ -1,4 +1,4 @@
-Installation
+Installation (Torch and JAX)
 ==============================
 
 .. toctree::
@@ -8,11 +8,15 @@ Installation
 You need the following to install OpenEquivariance:
 
 - A Linux system equipped with an NVIDIA / AMD graphics card.
-- PyTorch >= 2.4 (>= 2.8 for AOTI and export).
+- Either PyTorch >= 2.4 (>= 2.8 for AOTI and export), or JAX with CUDA 12 support
+  or higher.
 - GCC 9+ and the CUDA / HIP toolkit. The command
   ``c++ --version`` should return >= 9.0; see below for details on 
   setting an alternate compiler.
 
+PyTorch
+------------------------------------------
+
 Installation is one easy command, followed by import verification: 
 
 .. code-block:: bash 
@@ -28,11 +32,8 @@ To get the nightly build, run
 
 .. code-block:: bash 
 
-    pip install git+https://github.com/PASSIONLab/OpenEquivariance 
-
+    pip install git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance
 
-Compiling the Integrated PyTorch Extension
-------------------------------------------
 To support ``torch.compile``, ``torch.export``, and
 JITScript, OpenEquivariance needs to compile a C++ extension
 tightly integrated with PyTorch. If you see a warning that
@@ -48,13 +49,37 @@ environment variable and retry the import:
 
 .. code-block:: bash
 
-    export CCC=/path/to/your/gcc
+    export CC=/path/to/your/gcc
     export CXX=/path/to/your/g++
     python -c "import openequivariance"
 
 These configuration steps are required only ONCE after 
 installation (or upgrade) with pip. 
 
+JAX
+------------------------------------------
+JAX support is currently limited to NVIDIA GPUs. You need to execute
+the following two commands strictly in order:
+
+.. code-block:: bash 
+
+    pip install openequivariance[jax]
+    pip install openequivariance_extjax --no-build-isolation
+
+From there, set ``OEQ_NOTORCH=1`` to avoid a PyTorch import and test the package:
+
+.. code-block:: bash
+
+    OEQ_NOTORCH=1
+    python -c "import openequivariance.jax"
+
+You can get the nightly build as follows:
+
+.. code-block:: bash 
+
+    pip install git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance[jax]
+    pip install git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax
+
 Configurations on Major Platforms 
 ---------------------------------
 OpenEquivariance has been tested on both supercomputers and lab clusters.

openequivariance/openequivariance/__init__.py

@@ -31,21 +31,18 @@ def _check_package_editable():
 
 _editable_install_output_path = Path(__file__).parent.parent.parent / "outputs"
 
-
 def extension_source_path():
     """
     :returns: Path to the source code of the C++ extension.
     """
     return str(Path(__file__).parent / "extension")
 
-TensorProduct, TensorProductConv, torch_ext_so_path, torch_to_oeq_dtype = None, None, None, None
-
 if "OEQ_NOTORCH" not in os.environ or os.environ["OEQ_NOTORCH"] != "1":
     import torch
     from openequivariance.impl_torch.TensorProduct import TensorProduct 
     from openequivariance.impl_torch.TensorProductConv import TensorProductConv
 
-    from openequivariance.impl_torch.extlib import torch_ext_so_path
+    from openequivariance.impl_torch.extlib import torch_ext_so_path as torch_ext_so_path_internal
     from openequivariance.core.utils import torch_to_oeq_dtype
 
     torch.serialization.add_safe_globals(
@@ -62,6 +59,16 @@ def extension_source_path():
         ]
     )
 
+def torch_ext_so_path():
+    """
+    :returns: Path to a ``.so`` file that must be linked to use OpenEquivariance
+              from the PyTorch C++ Interface.
+    """
+    try:
+        return torch_ext_so_path_internal()
+    except NameError:
+        return None
+
 jax = None
 try:
     import openequivariance_extjax

openequivariance/openequivariance/impl_torch/extlib/__init__.py

@@ -142,10 +142,6 @@ def _raise_import_error_helper(import_target: str):
 
 
 def torch_ext_so_path():
-    """
-    :returns: Path to a ``.so`` file that must be linked to use OpenEquivariance
-              from the PyTorch C++ Interface.
-    """
     return torch_module.__file__