diff --git a/docs/conf.py b/docs/conf.py
index f5b9727..e76d6ce 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -21,6 +21,9 @@
 import sys
 
 
+# Add the project root directory to the Python path
+sys.path.insert(0, os.path.abspath(".."))
+
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
@@ -33,18 +36,46 @@
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
 extensions = [
-    "myst_parser",  # For our markdown docs
-    "autodoc2",  # Generates API docs
     "sphinx.ext.viewcode",  # For adding a link to view source code in docs
     "sphinx.ext.doctest",  # Allows testing in docstrings
     "sphinx.ext.napoleon",  # For google style docstrings
     "sphinx_copybutton",  # For copy button in code blocks
     "sphinxcontrib.katex",  # For KaTeX math rendering
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.intersphinx",
+    "sphinx.ext.todo",
+    "sphinx.ext.coverage",
+    "sphinx.ext.autosectionlabel",
+    "sphinx_panels",
+    "sphinx.ext.linkcode",
 ]
 
 templates_path = ["_templates"]
 exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
+autosummary_generate = True
+numpydoc_show_class_members = False
+panels_add_bootstrap_css = False
+autosectionlabel_prefix_document = True
+katex_prerender = False
+napoleon_use_ivar = True
+html_domain_indices = False
+source_suffix = ".rst"
+master_doc = "index"
+autodoc_docstring_signature = True
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3", None),
+    "numpy": ("https://numpy.org/doc/stable", None),
+    "torch": ("https://pytorch.org/docs/2.5", None),
+}
+
+
+def linkcode_resolve(domain, info):
+    """Resolve the linkcode for the given domain and info."""
+    return None
+
+
 # -- Options for MyST Parser (Markdown) --------------------------------------
 # MyST Parser settings
 myst_enable_extensions = [
@@ -57,21 +88,6 @@
 ]
 myst_heading_anchors = 5  # Generates anchor links for headings up to level 5
 
-# -- Options for Autodoc2 ---------------------------------------------------
-sys.path.insert(0, os.path.abspath(".."))
-
-autodoc2_packages = [
-    "../emerging_optimizers",  # Path to your package relative to conf.py
-]
-autodoc2_render_plugin = "myst"  # Use MyST for rendering docstrings
-autodoc2_output_dir = "apidocs"  # Output directory for autodoc2 (relative to docs/)
-# This is a workaround that uses the parser located in autodoc2_docstrings_parser.py to allow autodoc2 to
-# render google style docstrings.
-# Related Issue: https://github.com/sphinx-extensions2/sphinx-autodoc2/issues/33
-autodoc2_docstring_parser_regexes = [
-    (r".*", "docs.autodoc2_docstrings_parser"),
-]
-
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
diff --git a/docs/index.md b/docs/index.md
deleted file mode 100644
index 15f004a..0000000
--- a/docs/index.md
+++ /dev/null
@@ -1,11 +0,0 @@
-```{include} ../README.md
-:relative-docs: docs/
-```
-
-```{toctree}
-:caption: Development
-:hidden:
-
-documentation.md
-apidocs/index.rst
-```
diff --git a/docs/index.rst b/docs/index.rst
new file mode 100644
index 0000000..89566c9
--- /dev/null
+++ b/docs/index.rst
@@ -0,0 +1,16 @@
+Welcome to Emerging Optimizers's documentation!
+===============================================
+
+.. toctree::
+   :maxdepth: 1
+   :caption: API documentation
+
+   utils
+   orthogonalized_optimizers
+
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+* :ref:`search`
diff --git a/docs/orthogonalized_optimizers.rst b/docs/orthogonalized_optimizers.rst
new file mode 100644
index 0000000..44fe942
--- /dev/null
+++ b/docs/orthogonalized_optimizers.rst
@@ -0,0 +1,30 @@
+.. role:: hidden
+    :class: hidden-section
+
+emerging_optimizers.orthogonalized_optimizers
+=============================================
+
+.. automodule:: emerging_optimizers.orthogonalized_optimizers
+.. currentmodule:: emerging_optimizers.orthogonalized_optimizers
+
+:hidden:`OrthogonalizedOptimizer`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: OrthogonalizedOptimizer
+    :members:
+
+:hidden:`Muon`
+~~~~~~~~~~~~~~~
+
+.. autoclass:: Muon
+    :members:
+
+
+:hidden:`Newton-Schulz`
+~~~~~~~~~~~~~~~~~~~~~~~~
+.. automodule:: emerging_optimizers.orthogonalized_optimizers.muon_utils
+.. currentmodule:: emerging_optimizers.orthogonalized_optimizers.muon_utils
+
+.. autofunction:: newton_schulz
+.. autofunction:: newton_schulz_step
+.. autofunction:: newton_schulz_tp
diff --git a/docs/utils.rst b/docs/utils.rst
new file mode 100644
index 0000000..8bb16fb
--- /dev/null
+++ b/docs/utils.rst
@@ -0,0 +1,23 @@
+.. role:: hidden
+    :class: hidden-section
+
+emerging_optimizers.utils
+==========================
+
+.. automodule:: emerging_optimizers.utils
+.. currentmodule:: emerging_optimizers.utils
+
+:hidden:`fp32_matmul_precision`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: fp32_matmul_precision
+
+:hidden:`get_pg_size`
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: get_pg_size
+
+:hidden:`get_pg_rank`
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autofunction:: get_pg_rank
diff --git a/pyproject.toml b/pyproject.toml
index 52a4206..d29385e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -72,6 +72,7 @@ docs = [
     "sphinx-autobuild>=2024.10.3",
     "sphinx-autodoc2>=0.5.0",
     "sphinx-copybutton>=0.5.2",
+    "sphinx-panels>=0.4.1",
     "sphinxcontrib-katex>=0.9.11",
 ]
 test = [
diff --git a/uv.lock b/uv.lock
index a665dce..7c50205 100644
--- a/uv.lock
+++ b/uv.lock
@@ -325,6 +325,7 @@ docs = [
     { name = "sphinx-autobuild", version = "2025.8.25", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
     { name = "sphinx-autodoc2" },
     { name = "sphinx-copybutton" },
+    { name = "sphinx-panels" },
     { name = "sphinxcontrib-katex" },
 ]
 test = [
@@ -356,6 +357,7 @@ docs = [
     { name = "sphinx-autobuild", specifier = ">=2024.10.3" },
     { name = "sphinx-autodoc2", specifier = ">=0.5.0" },
     { name = "sphinx-copybutton", specifier = ">=0.5.2" },
+    { name = "sphinx-panels", specifier = ">=0.4.1" },
     { name = "sphinxcontrib-katex", specifier = ">=0.9.11" },
 ]
 test = [
@@ -1301,6 +1303,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/9e/48/1ea60e74949eecb12cdd6ac43987f9fd331156388dcc2319b45e2ebb81bf/sphinx_copybutton-0.5.2-py3-none-any.whl", hash = "sha256:fb543fd386d917746c9a2c50360c7905b605726b9355cd26e9974857afeae06e", size = 13343, upload-time = "2023-04-14T08:10:20.844Z" },
 ]
 
+[[package]]
+name = "sphinx-panels"
+version = "0.4.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "docutils" },
+    { name = "sphinx", version = "8.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "sphinx", version = "8.2.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b8/46/ca13f9739ba9783665be13ae8c066793d50ef50d75654160fd26c2812553/sphinx-panels-0.4.1.tar.gz", hash = "sha256:dc3c0add09b4264d3328d1b90857d429d9af65ae356d3ff7e437901946077533", size = 72668, upload-time = "2020-06-30T01:18:07.99Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/74/27/7f145fcbdd6d53990df2d4d22a9766bd47fca2ddf58a77c6476a9b9d4653/sphinx_panels-0.4.1-py3-none-any.whl", hash = "sha256:21c86a21f1109b00be87be9e057882ac8ca559d83805e58145a9f9a90f41d4df", size = 74331, upload-time = "2020-06-30T01:18:06.78Z" },
+]
+
 [[package]]
 name = "sphinxcontrib-applehelp"
 version = "2.0.0"