Source code for pymc_extras.inference.fit

#   Copyright 2022 The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
from xarray import DataTree




def fit(method: str, **kwargs) -> DataTree:
    """
    Fit a model with an inference algorithm.
    See :func:`fit_pathfinder` and :func:`fit_laplace` for more details.

    Parameters
    ----------
    method : str
        Which inference method to run.
        Supported: pathfinder or laplace

    kwargs: keyword arguments are passed on to the inference method.

    Returns
    -------
    DataTree
    """
    if method == "pathfinder":
        from pymc_extras.inference.pathfinder import fit_pathfinder

        return fit_pathfinder(**kwargs)

    if method == "laplace":
        from pymc_extras.inference import fit_laplace

        return fit_laplace(**kwargs)

    if method == "dadvi":
        from pymc_extras.inference import fit_dadvi

        return fit_dadvi(**kwargs)