diff --git a/dolo/algos/ergodic.py b/dolo/algos/ergodic.py
index 4f050382..bff9210b 100644
--- a/dolo/algos/ergodic.py
+++ b/dolo/algos/ergodic.py
@@ -1,8 +1,15 @@
 import numpy as np
+
+# doesn't seem to support keyword arguments yet
+# from multipledispatch import dispatch
+# multimethod = dispatch()
 from multimethod import multimethod
 from numba import generated_jit
 import xarray
 
+from dolo.compiler.model import Model
+from dolo.numeric.grids import UniformCartesianGrid
+from dolo.numeric.decision_rule import DecisionRule
 from dolo.numeric.processes import MarkovChain, IIDProcess, DiscretizedIIDProcess
 from dolo.numeric.grids import CartesianGrid, UnstructuredGrid, EmptyGrid
 
@@ -74,13 +81,13 @@ def fun(A: 'N* *...', x: 'N*d', w: 'float'):
 
 # TODO: add default options for endo_grid, exo_grid, dp
 @multimethod
-def ergodic_distribution(model, dr):
+def ergodic_distribution(model: Model, dr: DecisionRule):
     return ergodic_distribution(model, dr, dr.exo_grid, dr.endo_grid, dr.dprocess)
 
 @multimethod
-def ergodic_distribution(model, dr,
+def ergodic_distribution(model: Model, dr: DecisionRule,
         exo_grid: UnstructuredGrid,
-        endo_grid: CartesianGrid,
+        endo_grid: UniformCartesianGrid,
         dp: MarkovChain,
         compute_μ=True):
 
@@ -134,7 +141,7 @@ def ergodic_distribution(model, dr,
 
 
 @multimethod
-def ergodic_distribution(model, dr,
+def ergodic_distribution(model: Model, dr: DecisionRule,
         exo_grid: EmptyGrid,
         endo_grid: CartesianGrid,
         dp: DiscretizedIIDProcess,
diff --git a/environment.yml b/environment.yml
index 9bc9a116..e6493e77 100644
--- a/environment.yml
+++ b/environment.yml
@@ -18,5 +18,6 @@ dependencies:
 - matplotlib
 - xarray
 - multipledispatch
+- multimethod
 # - pip:
 #   - sphinx-better-theme==0.13
diff --git a/requirements.txt b/requirements.txt
index 1ccc002a..c1c4cf8e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 dolang
 multipledispatch
+multimethod
diff --git a/setup.py b/setup.py
index 3aa8a861..e41a1576 100644
--- a/setup.py
+++ b/setup.py
@@ -15,7 +15,7 @@
     ],
     install_requires=[
         "dolang", "pyyaml", "numba", "numpy", "sympy", "scipy", "quantecon", "pandas",
-        "interpolation", "ruamel.yaml", "xarray", "altair", "multipledispatch"
+        "interpolation", "ruamel.yaml", "xarray", "altair", "multipledispatch", "multimethod"
     ],
     extras_require={
         'interactive': ['ipython'],