WIP: Creation of Python autodoc using Sphinx

.github/workflows/sphinx.yml

@@ -0,0 +1,41 @@
+name: "Sphinx: Render docs"
+
+on: 
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build-docs:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.11'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install sphinx>=5.0.0
+        # Install other dependencies if needed
+        pip install -r requirements.txt
+
+    - name: Build documentation
+      run: |
+        cd docs
+        make html
+    - name: Upload artifacts
+      uses: actions/upload-artifact@v3
+      with:
+        name: documentation
+        path: docs/_build/html/
+
+    - name: Deploy to GitHub Pages
+      if: github.event_name == 'push' && github.ref == 'refs/heads/main'
+      uses: peaceiris/actions-gh-pages@v3
+      with:
+        github_token: ${{ secrets.GITHUB_TOKEN }}
+        publish_dir: ./docs/_build/html

docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/make.bat

@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd

docs/source/api.rst

@@ -0,0 +1,106 @@
+Liger Kernel API
+================
+
+.. automodule:: liger_kernel
+   :members:
+
+Submodules
+----------
+.. toctree::
+   :maxdepth: 2
+
+   liger_kernel.transformers
+   liger_kernel.nn
+
+Transformers Module
+-------------------
+
+The `liger_kernel.transformers` module includes essential tools and model classes designed for tasks such as language modeling and loss calculations, optimizing deep learning workflows.
+
+.. automodule:: liger_kernel.transformers
+   :members:
+
+Classes
+-------
+
+AutoLigerKernelForCausalLM
+--------------------------
+
+The `AutoLigerKernelForCausalLM` class is an extension for causal language modeling, providing enhanced support for auto-regressive tasks.
+
+.. autoclass:: liger_kernel.transformers.AutoLigerKernelForCausalLM
+   :members:
+   :undoc-members:
+
+LigerFusedLinearCrossEntropyLoss
+--------------------------------
+
+Implements a fused linear layer with Cross-Entropy Loss, offering computational efficiency improvements.
+
+.. autoclass:: liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss
+   :members:
+   :undoc-members:
+
+Loss Functions
+--------------
+
+KLDivergence
+------------
+
+The `KLDivergence` class offers a divergence metric based on Kullback-Leibler Divergence, commonly used in various ML training contexts.
+
+.. autoclass:: liger_kernel.transformers.KLDivergence
+   :members:
+   :undoc-members:
+
+JSD
+---
+
+This class defines the Jensen-Shannon Divergence (JSD), used for calculating symmetrical divergence between distributions.
+
+.. autoclass:: liger_kernel.transformers.JSD
+   :members:
+   :undoc-members:
+
+GeneralizedJSD
+--------------
+
+Generalized form of Jensen-Shannon Divergence, adapted for more complex applications.
+
+.. autoclass:: liger_kernel.transformers.GeneralizedJSD
+   :members:
+   :undoc-members:
+
+FusedLinearJSD
+--------------
+
+Provides a fused linear implementation of Jensen-Shannon Divergence for faster computation.
+
+.. autoclass:: liger_kernel.transformers.FusedLinearJSD
+   :members:
+   :undoc-members:
+
+Experimental Kernels
+---------------------
+
+A collection of experimental kernels for advanced and experimental features in `liger_kernel`.
+
+.. automodule:: liger_kernel.transformers.experimental
+   :members:
+
+Neural Network Module
+---------------------
+
+The `liger_kernel.nn` module offers fundamental neural network building blocks.
+
+.. automodule:: liger_kernel.nn
+   :members:
+
+Module Class
+------------
+
+The base `nn.Module` class provides essential methods and attributes for building neural network models.
+
+.. autoclass:: liger_kernel.nn.Module
+   :members:
+   :undoc-members:

docs/source/conf.py

@@ -0,0 +1,126 @@
+import os
+import sys
+import importlib
+import inspect
+
+# Increase the recursion limit (optional)
+sys.setrecursionlimit(10000)
+
+# Configure paths
+docs_dir = os.path.abspath(os.path.dirname(__file__))
+root_dir = os.path.abspath(os.path.join(docs_dir, '..'))
+sys.path.insert(0, root_dir)
+
+# Source and build directories
+source_dir = os.path.join(docs_dir, 'source')
+output_dir = os.path.join(docs_dir, '_build')
+
+# List modules to mock (as simple strings)
+MOCK_MODULES = [
+    'liger_kernel',
+    'liger_kernel.transformers',
+    'liger_kernel.transformers.experimental',
+    'liger_kernel.nn',
+    'liger_kernel.transformers.AutoLigerKernelForCausalLM',
+    'liger_kernel.transformers.LigerFusedLinearCrossEntropyLoss',
+    'liger_kernel.transformers.KLDivergence',
+    'liger_kernel.transformers.JSD',
+    'liger_kernel.transformers.GeneralizedJSD',
+    'liger_kernel.transformers.FusedLinearJSD',
+    'liger_kernel.nn.Module',
+]
+
+# Mock modules by replacing them in sys.modules
+for mod_name in MOCK_MODULES:
+    sys.modules[mod_name] = None  # Simple mock; replace with 'None' or a basic object as needed
+
+# Project information
+project = 'Liger-Kernel'
+copyright = '2024'
+author = 'LinkedIn'
+
+# General Sphinx configuration
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx.ext.napoleon',
+    'sphinx.ext.viewcode',
+    'sphinx.ext.todo',
+    'sphinx.ext.autosummary',
+    'sphinx.ext.linkcode',
+    'sphinx.ext.intersphinx',
+]
+
+# Mocked imports for autodoc
+autodoc_mock_imports = ['numpy', 'tensorflow', 'torch', 'triton']
+
+# Intersphinx mapping for cross-referencing external documentation
+intersphinx_mapping = {
+    'python': ('https://docs.python.org/3', None),
+    'numpy': ('https://numpy.org/doc/stable', None),
+    'torch': ('https://pytorch.org/docs/stable', None),
+    'triton': ('https://triton-lang.org/main/', None),
+}
+
+# Source file configuration
+source_suffix = '.rst'
+master_doc = 'index'
+exclude_patterns = []
+
+# HTML output configuration
+html_theme = 'alabaster'
+html_static_path = [os.path.join(source_dir, '_static')]
+html_theme_options = {
+    'navigation_depth': 4,
+    'collapse_navigation': False,
+}
+
+# Output directory for HTML files
+html_output_dir = output_dir
+
+# Autodoc configuration
+autodoc_member_order = 'bysource'
+autoclass_content = 'both'
+
+# Enable todo extension
+todo_include_todos = True
+
+# Autosummary settings
+autosummary_generate = True
+
+def linkcode_resolve(domain, info):
+    """
+    Determine the URL corresponding to Python object
+    """
+    if domain != 'py':
+        return None
+
+    try:
+        mod = importlib.import_module(info['module'])
+        if 'class' in info:
+            obj = getattr(mod, info['class'])
+            if 'fullname' in info:
+                obj = getattr(obj, info['fullname'])
+        else:
+            obj = getattr(mod, info['fullname'])
+    except Exception:
+        return None
+
+    try:
+        filepath = inspect.getsourcefile(obj)
+        if filepath:
+            filepath = os.path.relpath(filepath, start=root_dir)
+        else:
+            return None
+    except Exception:
+        return None
+
+    try:
+        source, lineno = inspect.getsourcelines(obj)
+    except Exception:
+        return None
+
+    # Modify these values based on your repository
+    github_url = "https://github.com/LinkedIn/liger-kernel"
+    branch = "main"  # or whatever your default branch is
+
+    return f"{github_url}/blob/{branch}/{filepath}#L{lineno}-L{lineno + len(source) - 1}"

docs/source/index.rst

@@ -0,0 +1,17 @@
+Welcome to Liger-Kernel Documentation
+===================================
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Contents:
+
+   source/api
+   source/modules
+   source/todo
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+* :ref:`modindex`
+* :ref:`search`

docs/source/liger_kernel.transformers.rst

@@ -0,0 +1,36 @@
+.. automodule:: liger_kernel.transformers
+   :members:
+
+Liger Cross Entropy Function
+=============================
+
+.. autoclass:: LigerCrossEntropyFunction
+   :members:
+
+   The Liger Cross Entropy Function is a custom autograd function that implements the Liger Cross Entropy loss.
+   It overrides the forward and backward methods of the torch.autograd.Function class.
+
+  .. automethod:: forward
+
+  .. automethod:: backward
+
+Liger Cross Entropy Forward
+============================
+
+.. autofunction:: cross_entropy_forward
+
+Liger Cross Entropy Backward
+============================
+
+.. autofunction:: cross_entropy_backward
+
+Liger Cross Entropy Kernel
+==========================
+
+.. autofunction:: liger_cross_entropy_kernel
+
+Element Mul Kernel
+===================
+
+.. autofunction:: element_mul_kernel
+