docs writing

Author: stes

cebra/models/__init__.py

@@ -36,7 +36,7 @@
 from cebra.models.multiobjective import *
 from cebra.models.layers import *
 from cebra.models.criterions import *
-from cebra.models.multi_criterions import *
+from cebra.models.multicriterions import *
 from cebra.models.jacobian_regularizer import *
 
 cebra.registry.add_docstring(__name__)

cebra/models/jacobian_regularizer.py

@@ -37,6 +37,14 @@
 #   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 #   SOFTWARE.
 #
+"""Jacobian Regularization for CEBRA.
+
+This implementation is adapted from the Jacobian regularization described in [1]_.
+
+.. [1] Judy Hoffman, Daniel A. Roberts, and Sho Yaida,
+       "Robust Learning with Jacobian Regularization," 2019.
+       `arxiv:1908.02729 <https://arxiv.org/abs/1908.02729>`_
+"""
 
 from __future__ import division
 
@@ -52,12 +60,6 @@ class JacobianReg(nn.Module):
            of the output space and projection is non-random and orthonormal, yielding the exact
            result. For any reasonable batch size, the default (n=1) should be sufficient.
            |Default:| ``1``
-
-    Note:
-        This implementation is adapted from the Jacobian regularization described in [1].
-        [1] Judy Hoffman, Daniel A. Roberts, and Sho Yaida,
-            "Robust Learning with Jacobian Regularization," 2019.
-            [arxiv:1908.02729](https://arxiv.org/abs/1908.02729)
     """
 
     def __init__(self, n: int = 1):
@@ -66,7 +68,7 @@ def __init__(self, n: int = 1):
         super(JacobianReg, self).__init__()
 
     def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        """Computes (1/2) tr |dy/dx|^2.
+        """Computes (1/2) tr \\|dy/dx\\|^2.
 
         Args:
             x: Input tensor

cebra/models/multi_criterions.py cebra/models/multicriterions.pycebra/models/multi_criterions.py renamed to cebra/models/multicriterions.py

@@ -19,6 +19,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+"""Support for training CEBRA with multiple criteria.
+
+.. note::
+   This module was introduced in CEBRA 0.6.0.
+
+"""
 from typing import Tuple
 
 import torch

docs/source/api.rst

@@ -38,7 +38,9 @@ these components in other contexts and research code bases.
    api/pytorch/distributions
    api/pytorch/models
    api/pytorch/helpers
-   api/pytorch/attribution
+   api/xcebra/multiobjective
+   api/xcebra/regularized
+   api/xcebra/attribution
 
 .. toctree::
    :hidden:

docs/source/api/pytorch/attribution.rst

@@ -42,13 +42,5 @@ Layers and model building blocks
    :members:
    :show-inheritance:
 
-Multi-objective models
-~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. automodule:: cebra.models.multiobjective
-   :members:
-   :private-members:
-   :show-inheritance:
-
 ..
    - projector

docs/source/api/pytorch/models.rst

@@ -0,0 +1,21 @@
+===================
+Attribution Methods
+===================
+
+.. automodule:: cebra.attribution
+   :members:
+   :show-inheritance:
+
+Different attribution methods
+-----------------------------
+
+.. automodule:: cebra.attribution.attribution_models
+   :members:
+   :show-inheritance:
+
+Jacobian-based attribution
+--------------------------
+
+.. automodule:: cebra.attribution.jacobian_attribution
+   :members:
+   :show-inheritance:

docs/source/api/xcebra/attribution.rst

@@ -0,0 +1,21 @@
+======================
+Multi-objective models
+======================
+
+Starting in CEBRA 0.6.0, we have added support for subspace contrastive learning.
+This is a method for training models that are able to learn multiple subspaces of the
+feature space simultaneously.
+
+Subspace contrastive learning requires to use specialized models and criterions.
+
+.. automodule:: cebra.models.multicriterions
+   :members:
+   :show-inheritance:
+
+.. automodule:: cebra.models.multiobjective
+   :members:
+   :show-inheritance:
+
+.. automodule:: cebra.solver.multiobjective
+   :members:
+   :show-inheritance:

docs/source/api/xcebra/multiobjective.rst

@@ -0,0 +1,24 @@
+================================
+Regularized Contrastive Learning
+================================
+
+Regularized solvers
+--------------------
+
+.. automodule:: cebra.solver.regularized
+   :members:
+   :show-inheritance:
+
+Schedulers
+----------
+
+.. automodule:: cebra.solver.schedulers
+   :members:
+   :show-inheritance:
+
+Jacobian Regularization
+-----------------------
+
+.. automodule:: cebra.models.jacobian_regularizer
+   :members:
+   :show-inheritance:

docs/source/api/xcebra/regularized.rst

@@ -203,7 +203,7 @@ def get_years(start_year=2021):
     ],
     "collapse_navigation": False,
     "navigation_depth": 1,
-    "show_nav_level": 2,
+    "show_nav_level": 1,
     "navbar_align": "content",
     "show_prev_next": False,
     "navbar_end": ["theme-switcher", "navbar-icon-links.html"],