Fix mkdocs errors; streamlining docs deploy with uv

.github/workflows/deploy-mkdocs.yml

@@ -11,19 +11,20 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+      - name: Set up uv
+        run: curl -LsSf https://astral.sh/uv/0.5.4/install.sh | sh
       - name: Configure Git Credentials
         run: |
           git config user.name github-actions[bot]
           git config user.email 41898282+github-actions[bot]@users.noreply.github.com
       - uses: actions/setup-python@v5
         with:
-          python-version: 3.x
+          python-version: 3.9
       - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV
       - uses: actions/cache@v4
         with:
           key: mkdocs-material-${{ env.cache_id }}
           path: .cache
           restore-keys: |
             mkdocs-material-
-      - run: pip install mkdocs-material mkdocstrings-python
-      - run: mkdocs gh-deploy --force
+      - run: uv run --only-group docs mkdocs gh-deploy --force

.github/workflows/pytest-asdfghjkl.yml

@@ -8,7 +8,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Set up uv
-        run: curl -LsSf https://astral.sh/uv/0.3.0/install.sh | sh
+        run: curl -LsSf https://astral.sh/uv/0.5.4/install.sh | sh
       - name: Set up Python
         run: uv python install 3.9
       - name: Install the project

.github/workflows/pytest-default.yml

@@ -8,7 +8,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Set up uv
-        run: curl -LsSf https://astral.sh/uv/0.3.0/install.sh | sh
+        run: curl -LsSf https://astral.sh/uv/0.5.4/install.sh | sh
       - name: Set up Python
         run: uv python install 3.9
       - name: Install the project

laplace/baselaplace.py

@@ -2102,14 +2102,48 @@ class FunctionalLaplace(BaseLaplace):
 
     Parameters
     ----------
-    num_data : int
+    model : torch.nn.Module
+    likelihood : Likelihood or str in {'classification', 'regression', 'reward_modeling'}
+        determines the log likelihood Hessian approximation.
+        In the case of 'reward_modeling', it fits Laplace using the classification likelihood,
+        then does prediction as in regression likelihood. The model needs to be defined accordingly:
+        The forward pass during training takes `x.shape == (batch_size, 2, dim)` with
+        `y.shape = (batch_size,)`. Meanwhile, during evaluation `x.shape == (batch_size, dim)`.
+        Note that 'reward_modeling' only supports `KronLaplace` and `DiagLaplace`.
+    sigma_noise : torch.Tensor or float, default=1
+        observation noise for the regression setting; must be 1 for classification
+    prior_precision : torch.Tensor or float, default=1
+        prior precision of a Gaussian prior (= weight decay);
+        can be scalar, per-layer, or diagonal in the most general case
+    prior_mean : torch.Tensor or float, default=0
+        prior mean of a Gaussian prior, useful for continual learning
+    temperature : float, default=1
+        temperature of the likelihood; lower temperature leads to more
+        concentrated posterior and vice versa.
+    enable_backprop: bool, default=False
+        whether to enable backprop to the input `x` through the Laplace predictive.
+        Useful for e.g. Bayesian optimization.
+    dict_key_x: str, default='input_ids'
+        The dictionary key under which the input tensor `x` is stored. Only has effect
+        when the model takes a `MutableMapping` as the input. Useful for Huggingface
+        LLM models.
+    dict_key_y: str, default='labels'
+        The dictionary key under which the target tensor `y` is stored. Only has effect
+        when the model takes a `MutableMapping` as the input. Useful for Huggingface
+        LLM models.
+    backend : subclasses of `laplace.curvature.CurvatureInterface`
+        backend for access to curvature/Hessian approximations. Defaults to CurvlinopsGGN if None.
+    backend_kwargs : dict, default=None
+        arguments passed to the backend on initialization, for example to
+        set the number of MC samples for stochastic approximations.
+    n_subset : int
         number of data points for Subset-of-Data (SOD) approximate GP inference.
-    diagonal_kernel : bool
+    independent_outputs : bool
         GP kernel here is product of Jacobians, which results in a \\( C \\times C\\) matrix where \\(C\\) is the output
         dimension. If `diagonal_kernel=True`, only a diagonal of a GP kernel is used. This is (somewhat) equivalent to
         assuming independent GPs across output channels.
-
-    See `BaseLaplace` class for the full interface.
+    seed: int, default=0
+        Random seed for subset of data sampler.
     """
 
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)

laplace/utils/matrix.py

@@ -297,7 +297,7 @@ class KronDecomposed:
     deltas : torch.Tensor
         addend for each group of Kronecker factors representing, for example,
         a prior precision
-    dampen : bool, default=False
+    damping : bool, default=False
         use dampen approximation mixing prior and Kron partially multiplicatively
     """
 

laplace/utils/subnetmask.py

@@ -356,7 +356,7 @@ class ModuleNameSubnetMask(SubnetMask):
     Parameters
     ----------
     model : torch.nn.Module
-    parameter_names: List[str]
+    module_names: List[str]
         list of names of the modules (as in `model.named_modules()`) that define the subnetwork;
         the modules cannot have children, i.e. need to be leaf modules
     """

pyproject.toml

@@ -42,18 +42,18 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools]
 py-modules = ["laplace"]
 
-[project.optional-dependencies]
-dev = [
-  "coveralls",
+[tool.uv]
+default-groups = ["dev", "docs"]
+
+[dependency-groups]
+dev = ["coveralls", "pytest", "pytest-cov", "pytest-mock", "ruff", "scipy"]
+docs = [
+  "black",
   "matplotlib",
-  "mkdocs",
-  "mkdocs-material",
-  "mkdocstrings-python",
-  "pytest",
-  "pytest-cov",
-  "pytest-mock",
-  "ruff",
-  "scipy",
+  "mkdocs == 1.6.1",
+  "mkdocs-material == 9.5.34",
+  "mkdocstrings == 0.26.1",
+  "mkdocstrings-python == 1.11.1",
 ]
 
 [tool.pytest.ini_options]