diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 556bf46..b090dc7 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -17,11 +17,11 @@ jobs:
       - uses: actions/setup-python@v4
 
       - name: Setup Miniconda
-        uses: conda-incubator/setup-miniconda@v2
+        uses: conda-incubator/setup-miniconda@v3
         with:
           auto-update-conda: true
           python-version: 3.8
-          miniforge-variant: Mambaforge
+          miniforge-variant: Miniforge3
           miniforge-version: latest
 
       - name: Conda info
@@ -30,7 +30,7 @@ jobs:
           conda info
           which python
 
-      - name: Mamba install FEniCS
+      - name: conda install FEniCS
         shell: bash -l {0}
         run: |
           conda config --set always_yes yes
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index da12601..a170e33 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -19,11 +19,11 @@ jobs:
           fetch-depth: 1
 
       - name: Setup Miniconda
-        uses: conda-incubator/setup-miniconda@v2
+        uses: conda-incubator/setup-miniconda@v3
         with:
           auto-update-conda: true
           python-version: 3.8
-          miniforge-variant: Mambaforge
+          miniforge-variant: Miniforge3
           miniforge-version: latest
 
       - name: Conda info
@@ -32,12 +32,12 @@ jobs:
           conda info
           which python
 
-      - name: Mamba install FEniCS
+      - name: Conda install FEniCS
         shell: bash -l {0}
         run: |
           conda config --set always_yes yes
           conda config --add channels conda-forge
-          mamba create -n fenicsproject -c conda-forge fenics 
+          conda create -n fenicsproject -c conda-forge fenics 
           conda activate fenicsproject
           which python
           python -c "from dolfin import *"
diff --git a/cuqipy_fenics/pde.py b/cuqipy_fenics/pde.py
index 757a9ff..359d09a 100644
--- a/cuqipy_fenics/pde.py
+++ b/cuqipy_fenics/pde.py
@@ -2,9 +2,11 @@
 from abc import ABC, abstractmethod
 from cuqi.pde import PDE
 from cuqi.array import CUQIarray
+from cuqi.utilities import get_non_default_args
 import dolfin as dl
 from copy import copy
 import warnings
+from functools import partial
 from .utilities import _LazyUFLLoader
 ufl = _LazyUFLLoader()
 
@@ -23,16 +25,18 @@ class FEniCSPDE(PDE,ABC):
     ----------
     PDE_form : callable or tuple of two callables
         If passed as a callable: the callable returns the weak form of the PDE.
-        The callable should take three arguments, the first argument is the 
-        parameter (input of the forward model), the second argument is the state
-        variable (solution variable), and the third argument is the adjoint
+        The callable should take three or more arguments, the first arguments
+        are the unknown parameters (inputs of the forward model, e.g.
+        `parameter1`, `parameter2`), the second to last argument is the state
+        variable (solution variable), and the last argument is the adjoint
         variable (the test variable in the weak formulation).
 
         If passed as a tuple of two callables: the first callable returns the 
         weak form of the PDE left hand side, and the second callable returns the
         weak form of the PDE right hand side. The left hand side callable takes
-        the same three arguments as described above. The right hand side 
-        callable takes only the parameter and the adjoint variable as arguments.
+        the same arguments as described above. The right hand side callable
+        takes only the unknown parameters and the adjoint variable (the latter 
+        being the last argument) as arguments.
         See the example below.
 
     mesh : FEniCS mesh
@@ -41,8 +45,9 @@ class FEniCSPDE(PDE,ABC):
     solution_function_space : FEniCS function space
         FEniCS function space object that defines the function space of the state variable (solution variable).
 
-    parameter_function_space : FEniCS function space
-        FEniCS function space object that defines the function space of the Bayesian parameter (input of the forward model).
+    parameter_function_space : FEniCS function space or a list of them
+        FEniCS function space object or a list of them that defines the function space of the unknown parameters (inputs of the forward model).
+        If multiple parameters are passed, the function space should be a list of FEniCS function spaces, one for each parameter.
 
     dirichlet_bcs: FEniCS Dirichlet boundary condition object or a list of them
         FEniCS Dirichlet boundary condition object(s) that define the Dirichlet boundary conditions of the PDE.
@@ -52,7 +57,7 @@ class FEniCSPDE(PDE,ABC):
 
     observation_operator : python function handle, optional
         Function handle of a python function that returns the observed quantity from the PDE solution. If not provided, the identity operator is assumed (i.e. the entire solution is observed).
-        This python function takes as input the Bayesian parameter (input of the forward model) and the state variable (solution variable) as first and second inputs, respectively.
+        This python function takes as input the unknown parameters, e.g. `parameter1`, `parameter2`, and the state variable (solution variable) in that order.
 
         The returned observed quantity can be a ufl.algebra.Operator, FEniCS Function, np.ndarray, int, or float.
 
@@ -146,9 +151,15 @@ def __init__(self, PDE_form, mesh, solution_function_space,
         for key, value in linalg_solve_kwargs.items():
             self._solver.parameters[key] = value
 
-        # Initialize the parameter
-        self.parameter = dl.Function(self.parameter_function_space)
-
+        # Initialize the parameter (one or more)
+        # If only one parameter is passed, it is converted to a list
+        if not isinstance(self.parameter_function_space, (list, tuple)):
+            parameter_function_space_list = [self.parameter_function_space]
+        else:
+            parameter_function_space_list = self.parameter_function_space
+        self.parameter= {}
+        for i, k in enumerate(self._non_default_args):
+            self.parameter[k] = dl.Function(parameter_function_space_list[i])
 
     @property
     def parameter(self):
@@ -157,8 +168,8 @@ def parameter(self):
     
     @parameter.setter
     def parameter(self, value):
-        """ Set the parameter of the PDE. Since the PDE solution depends on the 
-        parameter, this will set the PDE solution to None. """
+        """ Set the parameters of the PDE. Since the PDE solution depends on the 
+        parameters, this will set the PDE solution to None. """
         if value is None:
             raise ValueError('Parameter cannot be None.')
 
@@ -169,7 +180,8 @@ def parameter(self, value):
         # Subsequent times setting the parameter (avoid assigning the parameter
         # to new object, set parameter array in place instead)
         elif self._is_parameter_new(value):
-            self._parameter.vector().set_local(value.vector().get_local())
+            for key in self._non_default_args:
+                self._parameter[key].vector().set_local(value[key].vector().get_local())
             # The operator in the solver is no longer valid
             self._flags["is_operator_valid"] = False
 
@@ -179,6 +191,25 @@ def parameter(self, value):
         self._gradient = None
         self.rhs = None
 
+    @property
+    def parameter_args(self):
+        """Get the args form of the parameter"""
+        args = list(self.parameter.values())
+        return args
+
+    @property
+    def _non_default_args(self):
+        form = self._form
+        if isinstance(self._form, tuple):
+            # Use the lhs form only to determine the non-default args
+            form = self._form[0]
+
+        # Extract the non_default_args and exclude the last
+        # two arguments (u and p) from the list of non-default args since
+        # they are provided automatically within the PDE-type class and are
+        # not arguments to be inferred in Bayesian inference setting.
+        non_default_args = get_non_default_args(form)[:-2]
+        return non_default_args
 
     @property
     def forward_solution(self):
@@ -189,7 +220,24 @@ def forward_solution(self):
     def PDE_form(self):
         """ Get the PDE form """
         if isinstance(self._form, tuple):
-            return lambda m, u, p: self._form[0](m, u, p) - self._form[1](m, p)
+            # Create a string for the lambda function that represents the PDE form
+            form_str = (
+                "lambda form_lhs, form_rhs, "
+                + ", ".join(self._non_default_args)
+                + ", u, p: form_lhs("
+                + ", ".join(self._non_default_args)
+                + ", u, p) - form_rhs("
+                + ", ".join(self._non_default_args)
+                + ", p)"
+            )
+            # Create a lambda function that represents the PDE form
+            form = eval(form_str)
+
+            # partial evaluation of the form 
+            form_partial = partial(form, form_lhs=self._form[0],
+                                      form_rhs=self._form[1])
+
+            return form_partial
         else:
             return self._form
 
@@ -243,10 +291,14 @@ def observation_operator(self):
     @observation_operator.setter
     def observation_operator(self, value):
         """ Set the observation operator """
-        self._observation_operator = self._create_observation_operator(value)
+        if value == None or callable(value):
+            self._observation_operator = value
+        else:
+            raise NotImplementedError(
+                "observation_operator must be a callable function or None")
 
     @abstractmethod
-    def assemble(self,parameter):
+    def assemble(self, *args, **kwargs):
         """ Assemble the PDE weak form """
         raise NotImplementedError
 
@@ -265,34 +317,34 @@ def gradient_wrt_parameter(self):
         """ Compute gradient of the PDE weak form w.r.t. the parameter"""
         raise NotImplementedError 
 
-    @abstractmethod
-    def _create_observation_operator(self, observation_operator):
-        raise NotImplementedError
-
     def _is_parameter_new(self, input_parameter):
         """ A helper function to check if the `input_parameter` is different 
         from the current parameter (cached in self._parameter). """
 
-        if hasattr(self, '_parameter') \
-            and np.allclose(self._parameter.vector().get_local(),
-                            input_parameter.vector().get_local(),
-                            atol=dl.DOLFIN_EPS, rtol=dl.DOLFIN_EPS):
-            return False
-        else:
+        if not hasattr(self, '_parameter'):
             return True
 
+        is_new = False
+        for key in self._non_default_args:
+            if not np.allclose(self._parameter[key].vector().get_local(),
+                                    input_parameter[key].vector().get_local(),
+                                    atol=dl.DOLFIN_EPS, rtol=dl.DOLFIN_EPS):
+                is_new = True
+        return is_new
 
 class SteadyStateLinearFEniCSPDE(FEniCSPDE):
     """ Class representation of steady state linear PDEs defined in FEniCS. It accepts the same arguments as the base class `cuqipy_fenics.pde.FEniCSPDE`."""
 
-    def assemble(self, parameter=None):
+    def assemble(self, *args, **kwargs):
         self._solution_trial_function = dl.TrialFunction(
             self.solution_function_space)
         self._solution_test_function = dl.TestFunction(
             self.solution_function_space)
 
-        if parameter is not None:
-            self.parameter = parameter
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, map_name="assemble", **kwargs
+        )
+        self.parameter = kwargs
 
         # Either assemble the lhs and rhs forms separately or the full PDE form
         if self.lhs_form is not None:
@@ -328,11 +380,10 @@ def _assemble_full(self):
                 and self._flags["is_operator_valid"] and\
                 self.rhs is not None:
             return
-
-        diff_op = dl.lhs(self.PDE_form(self.parameter,
+        diff_op = dl.lhs(self.PDE_form(*self.parameter_args,
                                        self._solution_trial_function,
                                        self._solution_test_function))
-        self.rhs = dl.rhs(self.PDE_form(self.parameter,
+        self.rhs = dl.rhs(self.PDE_form(*self.parameter_args,
                                         self._solution_trial_function,
                                         self._solution_test_function))
 
@@ -353,7 +404,7 @@ def _assemble_lhs(self):
                 and self._flags["is_operator_valid"]:
             return
 
-        diff_op = dl.assemble(self.lhs_form(self.parameter,
+        diff_op = dl.assemble(self.lhs_form(*self.parameter_args,
                                             self._solution_trial_function,
                                             self._solution_test_function))
 
@@ -367,7 +418,7 @@ def _assemble_rhs(self):
                 and self.rhs is not None:
             return
 
-        self.rhs = dl.assemble(self.rhs_form(self.parameter,
+        self.rhs = dl.assemble(self.rhs_form(*self.parameter_args,
                                              self._solution_test_function))
         for bc in self._dirichlet_bcs: bc.apply(self.rhs)
 
@@ -381,16 +432,20 @@ def observe(self,PDE_solution_fun):
         if self.observation_operator is None: 
             return PDE_solution_fun
         else:
-            return self._apply_obs_op(self.parameter, PDE_solution_fun)
+            return self._apply_obs_op(*self.parameter_args, PDE_solution_fun)
 
-    def gradient_wrt_parameter(self, direction, wrt, **kwargs):
-        """ Compute the gradient of the PDE with respect to the parameter
+    def gradient_wrt_parameter(self, direction, *args, **kwargs):
+        """ Compute the gradient of the PDE with respect to the parameters
 
         Note: This implementation is largely based on the code:
         https://github.com/hippylib/hippylib/blob/master/hippylib/modeling/PDEProblem.py
 
         See also: Gunzburger, M. D. (2002). Perspectives in flow control and optimization. Society for Industrial and Applied Mathematics, for adjoint based derivative derivation. 
         """
+
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, map_name="gradient_wrt_parameter", **kwargs
+        )
         # Raise an error if the adjoint boundary conditions are not provided
         if self._adjoint_dirichlet_bcs is None:
             raise ValueError(
@@ -402,16 +457,15 @@ def gradient_wrt_parameter(self, direction, wrt, **kwargs):
 
         # Compute forward solution
         # TODO: Use stored forward solution if available and wrt == self.parameter
-        self.parameter = wrt
-        self.assemble()
+        self.parameter = kwargs
+        self.assemble(*self.parameter_args)
         self.forward_solution, _ = self.solve()
 
         # Compute adjoint solution
-        test_parameter = dl.TestFunction(self.parameter_function_space)
         test_solution = dl.TestFunction(self.solution_function_space)
 
         # note: temp_form is a weak form used for building the adjoint operator
-        temp_form = self.PDE_form(wrt, self.forward_solution, trial_adjoint)
+        temp_form = self.PDE_form(*self.parameter_args, self.forward_solution, trial_adjoint)
         adjoint_form = dl.derivative(
             temp_form, self.forward_solution, test_solution)
 
@@ -431,11 +485,25 @@ def gradient_wrt_parameter(self, direction, wrt, **kwargs):
 
         # Compute gradient
         # note: temp_form is a weak form used for building the gradient
-        temp_form = self.PDE_form(wrt, self.forward_solution, adjoint)
-        gradient_form = dl.derivative(temp_form, wrt, test_parameter)
-        gradient = dl.Function(self.parameter_function_space)
-        dl.assemble(gradient_form, tensor=gradient.vector())
-        return gradient
+        temp_form = self.PDE_form(*self.parameter_args, self.forward_solution, adjoint)
+        parameter_function_space = self.parameter_function_space
+        gradient_list = []
+        if not isinstance(self.parameter_function_space, (list, tuple)):
+            parameter_function_space = [self.parameter_function_space]
+
+        for i, k in enumerate(self._non_default_args):
+            test_parameter = dl.TestFunction(parameter_function_space[i])
+            gradient_form = dl.derivative(temp_form, self.parameter_args[i], test_parameter)
+            gradient = dl.Function(parameter_function_space[i])
+            dl.assemble(gradient_form, tensor=gradient.vector())
+            gradient_list.append(gradient)
+
+        # If only one parameter is passed, return a single gradient
+        if len(gradient_list) == 1:
+            gradient_list = gradient_list[0]
+        else:
+            gradient_list = tuple(gradient_list)
+        return gradient_list
 
 
     def _apply_obs_op(self, PDE_parameter_fun, PDE_solution_fun,):
@@ -450,21 +518,4 @@ def _apply_obs_op(self, PDE_parameter_fun, PDE_solution_fun,):
             raise NotImplementedError("obs_op output must be a number, a numpy array or a ufl.algebra.Operator type")
     
 
-    def _create_observation_operator(self, observation_operator):
-        """
-        """
-        if observation_operator == 'potential':
-            observation_operator = lambda m, u: u 
-        elif observation_operator == 'gradu_squared':
-            observation_operator = lambda m, u: dl.inner(dl.grad(u),dl.grad(u))
-        elif observation_operator == 'power_density':
-            observation_operator = lambda m, u: m*dl.inner(dl.grad(u),dl.grad(u))
-        elif observation_operator == 'sigma_u':
-            observation_operator = lambda m, u: m*u
-        elif observation_operator == 'sigma_norm_gradu':
-            observation_operator = lambda m, u: m*dl.sqrt(dl.inner(dl.grad(u),dl.grad(u)))
-        elif observation_operator == None or callable(observation_operator):
-            observation_operator = observation_operator
-        else:
-            raise NotImplementedError
-        return observation_operator
+
diff --git a/tests/test_PDEModel.py b/tests/test_PDEModel.py
index 910419f..3a8a35e 100644
--- a/tests/test_PDEModel.py
+++ b/tests/test_PDEModel.py
@@ -5,6 +5,7 @@
 import pytest
 import time
 import sys
+from scipy import optimize
 ufl = cuqipy_fenics.utilities._import_ufl()
 
 def test_model_input():
@@ -42,7 +43,7 @@ def test_model_input():
 def test_solver_choice():
     """Test passing different solvers to PDEModel"""
     # Create the variational problem
-    poisson = Poisson()
+    poisson = Poisson(dl.UnitIntervalMesh(10))
 
     # Set up model parameters
     m = dl.Function(poisson.parameter_function_space)
@@ -82,8 +83,7 @@ def test_reuse_assembled():
     """Test that reusing the assembled and factored lhs gives the same solution
      and a better performance"""
     # Create the variational problem
-    poisson = Poisson()
-    poisson.mesh = dl.UnitSquareMesh(50, 50)
+    poisson = Poisson(dl.UnitSquareMesh(50, 50))
 
     # Set up model parameters
     m = dl.Function(poisson.parameter_function_space)
@@ -139,7 +139,7 @@ def test_reuse_assembled():
 def test_form():
     """Test creating PDEModel with full form, and with lhs and rhs forms"""
     # Create the variational problem
-    poisson = Poisson()
+    poisson = Poisson(dl.UnitIntervalMesh(10))
 
     # Set up model parameters
     m = dl.Function(poisson.parameter_function_space)
@@ -184,24 +184,13 @@ def test_with_updated_rhs(copy_reference, case):
         pytest.skip("Test on MAC OS only")
 
     # Set up first poisson problem
-    poisson1 = Poisson()
-    poisson1.mesh = dl.UnitSquareMesh(40, 40)
+    poisson1 = Poisson(dl.UnitSquareMesh(40, 40))
     poisson1.source_term = dl.Constant(1.0)
 
     # Set up second poisson problem
-    poisson2 = Poisson()
-    poisson2.mesh = poisson1.mesh
+    poisson2 = Poisson(poisson1.mesh)
     poisson2.source_term = dl.Expression("sin(2*x[0]*pi)*sin(2*x[1]*pi)", degree=1)
 
-    # Set up boundary function (where the Dirichlet boundary conditions are applied)
-    def u_boundary(x, on_boundary):
-        return on_boundary and\
-            (x[0] < dl.DOLFIN_EPS or x[0] > 1.0 - dl.DOLFIN_EPS)
-
-    poisson1.bcs = dl.DirichletBC(poisson1.solution_function_space,
-                                 dl.Constant(0.0), u_boundary)
-    poisson2.bcs = poisson1.bcs
-
     # Create two PDE objects with different rhs terms
     PDE1 = cuqipy_fenics.pde.SteadyStateLinearFEniCSPDE(
         (poisson1.lhs_form, poisson1.rhs_form),
@@ -309,17 +298,26 @@ class Poisson:
     """Define the variational PDE problem for the Poisson equation in two
     ways: as a full form, and as lhs and rhs forms"""
 
-    def __init__(self):
+    def __init__(self, mesh):
 
-        # Create the mesh and define function spaces for the solution and the
-        # parameter
-        self.mesh = dl.UnitIntervalMesh(10)
+        # Set the mesh
+        self._mesh = mesh
 
         # Define the boundary condition
         self.bc_value = dl.Constant(0.0)
 
-        # the source term
-        self.source_term = dl.Expression('x[0]', degree=1)
+        # The source term
+        self._source_term = dl.Expression('x[0]', degree=1)
+
+        # Solution function space
+        self._solution_function_space = dl.FunctionSpace(self.mesh, "Lagrange", 2)
+
+        # Parameter function space
+        self._parameter_function_space = dl.FunctionSpace(self.mesh, "Lagrange", 1)
+
+    @property
+    def mesh(self):
+        return self._mesh
 
     @property
     def form(self):
@@ -338,29 +336,33 @@ def rhs_form(self):
 
     @property
     def solution_function_space(self):
-        return dl.FunctionSpace(self.mesh, "Lagrange", 2)
+        return self._solution_function_space
 
     @property
     def parameter_function_space(self):
-        return dl.FunctionSpace(self.mesh, "Lagrange", 1)
+        return self._parameter_function_space
 
     @property
     def bcs(self):
-        if not hasattr(self, "_bcs") or self._bcs is None:
-            self._bcs = dl.DirichletBC(self.solution_function_space,
-                                      self.bc_value, "on_boundary")
-        return self._bcs
-
-    @bcs.setter
-    def bcs(self, bcs):
-        self._bcs = bcs
+        return dl.DirichletBC(
+            self.solution_function_space, self.bc_value, "on_boundary"
+        )
+    
+    @property
+    def source_term(self):
+        """Return the source term"""
+        return self._source_term
+    
+    @source_term.setter
+    def source_term(self, value):
+        """Set the source term"""
+        self._source_term = value
 
 
 def test_observation_operator_setter():
     """Test that the observation setter works as expected"""
     # Create the variational problem
-    poisson = Poisson()
-    poisson.mesh = dl.UnitSquareMesh(50, 50)
+    poisson = Poisson(dl.UnitSquareMesh(50, 50))
 
     # Set up model parameters
     m = dl.Function(poisson.parameter_function_space)
@@ -404,3 +406,236 @@ def test_observation_operator_setter():
     # Check that the solutions are the same
     assert np.allclose(u1_obs, u2_obs) and np.allclose(u2_obs, u3_obs)
     assert len(u1_obs) == 5
+
+def test_gradient_poisson():
+    """Test the gradient of the Poisson PDEModel is computed correctly"""
+    # Set random seed
+    np.random.seed(0)
+
+    # Create the variational problem
+    poisson = Poisson(dl.UnitIntervalMesh(40))
+
+    # Set up model parameters
+    m = dl.Function(poisson.parameter_function_space)
+    m.vector()[:] = 1.0
+
+    # Create a PDE object
+    PDE = cuqipy_fenics.pde.SteadyStateLinearFEniCSPDE(
+        poisson.form,
+        poisson.mesh,
+        poisson.solution_function_space,
+        poisson.parameter_function_space,
+        poisson.bcs,
+        poisson.bcs)
+
+    # Create a PDE model
+    PDE_model = cuqi.model.PDEModel(
+        PDE,
+        domain_geometry=cuqipy_fenics.geometry.FEniCSContinuous(
+            poisson.parameter_function_space
+        ),
+        range_geometry=cuqipy_fenics.geometry.FEniCSContinuous(
+            poisson.solution_function_space
+        ),
+    )
+
+    # Create a prior distribution
+    m_prior = cuqi.distribution.Gaussian(
+        mean=np.zeros(PDE_model.domain_dim), cov=1, geometry=PDE_model.domain_geometry
+    )
+
+    # Create a likelihood
+    y = cuqi.distribution.Gaussian(
+        mean=PDE_model(m_prior), cov=np.ones(PDE_model.range_dim)*.1**2, geometry=PDE_model.range_geometry)
+    y = y(y=PDE_model(m.vector().get_local()))
+
+    # Evaluate the adjoint based gradient at value m2
+    m2 = dl.Function(poisson.parameter_function_space)
+    m2.vector()[:] = np.random.randn(PDE_model.domain_dim)
+    adjoint_grad = y.gradient(m_prior=m2.vector().get_local())
+
+    # Compute the FD gradient
+    step = 1e-7   # finite diff step
+    FD_grad = optimize.approx_fprime(m2.vector().get_local(), y.logd, step)
+
+    # Check that the adjoint gradient and FD gradient are close
+    assert np.allclose(adjoint_grad, FD_grad, rtol=1e-1),\
+        f"Adjoint gradient: {adjoint_grad}, FD gradient: {FD_grad}"
+    assert np.linalg.norm(adjoint_grad-FD_grad)/ np.linalg.norm(FD_grad) < 1e-2
+
+class PoissonMultipleInputs:
+    """Define the variational PDE problem for the Poisson equation with multiple
+     unknowns in two ways: as a full form, and as lhs and rhs forms"""
+
+    def __init__(self):
+
+        # Create the mesh and define function spaces for the solution and the
+        # parameter
+        self._mesh = dl.UnitIntervalMesh(10)
+
+        # Define the boundary condition
+        self.bc_value = dl.Constant(0.0)
+
+        # Set the solution function space
+        self._solution_function_space =  dl.FunctionSpace(self.mesh, "Lagrange", 2)
+
+        # Set the parameter function space
+        self._parameter_function_space = [
+            dl.FunctionSpace(self.mesh, "Lagrange", 1),
+            dl.FunctionSpace(self.mesh, "Lagrange", 1)]
+
+    @property
+    def mesh(self):
+        return self._mesh
+
+    @property
+    def form(self):
+        return lambda m, source_term, u, v:\
+            ufl.exp(m)*ufl.inner(ufl.grad(u), ufl.grad(v))*ufl.dx\
+            + source_term*v*ufl.dx
+
+    @property
+    def lhs_form(self):
+        return lambda m, source_term, u, v:\
+            ufl.exp(m)*ufl.inner(ufl.grad(u), ufl.grad(v))*ufl.dx
+
+    @property
+    def rhs_form(self):
+        return lambda m, source_term, v: -source_term*v*ufl.dx
+
+    @property
+    def solution_function_space(self):
+        return self._solution_function_space
+
+    @property
+    def parameter_function_space(self):
+        return self._parameter_function_space
+
+    @property
+    def bcs(self):
+        if not hasattr(self, "_bcs") or self._bcs is None:
+            self._bcs = dl.DirichletBC(self.solution_function_space,
+                                      self.bc_value, "on_boundary")
+        return self._bcs
+
+    @bcs.setter
+    def bcs(self, bcs):
+        self._bcs = bcs
+
+def test_form_multiple_inputs():
+    """Test creating PDEModel with full form, and with lhs and rhs forms with multiple inputs"""
+    # Create the variational problem
+    poisson = PoissonMultipleInputs()
+
+    # Set up model parameters
+    m = dl.Function(poisson.parameter_function_space[0])
+    m.vector()[:] = 1.0
+    source_term = dl.Function(poisson.parameter_function_space[1])
+    source_term.interpolate(dl.Expression('x[0]', degree=1))
+
+    # Create a PDE object with full form
+    PDE_with_full_form = cuqipy_fenics.pde.SteadyStateLinearFEniCSPDE(
+        poisson.form,
+        poisson.mesh,
+        poisson.solution_function_space,
+        poisson.parameter_function_space,
+        poisson.bcs)
+
+    # Solve the PDE
+    PDE_with_full_form.assemble(m, source_term)
+    # test also assembling with keyword arguments
+    PDE_with_full_form.assemble(m=m, source_term=source_term)    
+    u1, info = PDE_with_full_form.solve()
+
+    # Create a PDE object with lhs and rhs forms
+    PDE_with_lhs_rhs_forms = cuqipy_fenics.pde.SteadyStateLinearFEniCSPDE(
+        (poisson.lhs_form, poisson.rhs_form),
+        poisson.mesh,
+        poisson.solution_function_space,
+        poisson.parameter_function_space,
+        poisson.bcs)
+
+    # Solve the PDE
+    PDE_with_lhs_rhs_forms.assemble(m, source_term)
+    u2, info = PDE_with_lhs_rhs_forms.solve()
+
+    # Check that the solutions are the same
+    assert np.allclose(u1.vector().get_local(), u2.vector().get_local())
+
+    # Check passing wrong parameter_function_space raises an error
+    with pytest.raises(ValueError, match= r"assemble input is specified by keywords arguments \['m2', 'source_term'\] that does not match the non_default_args of assemble \['m', 'source_term'\]"):
+        PDE_with_full_form.assemble(m2=m, source_term=source_term)
+
+def test_gradient_poisson_multiple_inputs():
+    """Test the gradient of the Poisson PDEModel with multiple inputs is computed correctly"""
+    # Set random seed
+    np.random.seed(0)
+
+    # Create the variational problem
+    poisson = PoissonMultipleInputs()
+
+    # Set up model parameters
+    m = dl.Function(poisson.parameter_function_space[0])
+    m.vector()[:] = 1.0
+    source_term = dl.Function(poisson.parameter_function_space[1])
+    source_term.interpolate(dl.Expression('x[0]', degree=1))
+
+    # Create a PDE object
+    PDE = cuqipy_fenics.pde.SteadyStateLinearFEniCSPDE(
+        poisson.form,
+        poisson.mesh,
+        poisson.solution_function_space,
+        poisson.parameter_function_space,
+        poisson.bcs,
+        poisson.bcs)
+
+    domain_geom1 = cuqipy_fenics.geometry.FEniCSContinuous(
+                poisson.parameter_function_space[0]
+            )
+    domain_geom2 = cuqipy_fenics.geometry.FEniCSContinuous(
+                poisson.parameter_function_space[1]
+            )
+
+    # Create a PDE model
+    PDE_model = cuqi.model.PDEModel(
+        PDE,
+        domain_geometry=(domain_geom1, domain_geom2),
+        range_geometry=cuqipy_fenics.geometry.FEniCSContinuous(
+            poisson.solution_function_space)
+    )
+
+    # Create prior distributions
+    m_prior = cuqi.distribution.Gaussian(
+        mean=np.zeros(domain_geom1.par_dim), cov=1, geometry=domain_geom1
+    )
+    source_term_prior = cuqi.distribution.Gaussian(
+        mean=np.zeros(domain_geom2.par_dim), cov=1, geometry=domain_geom2
+    )
+
+    # Create a likelihood
+    y = cuqi.distribution.Gaussian(
+        mean=PDE_model(m_prior, source_term_prior), cov=np.ones(PDE_model.range_dim)*.1**2, geometry=PDE_model.range_geometry)
+    y = y(y=PDE_model(m.vector().get_local(), source_term.vector().get_local()))
+
+    # Evaluate the adjoint based gradient at value m2 and source_term2
+    m2 = dl.Function(poisson.parameter_function_space[0])
+    m2.vector()[:] = np.random.randn(domain_geom1.par_dim)
+    source_term2 = dl.Function(poisson.parameter_function_space[1])
+    source_term2.interpolate(dl.Expression('sin(2*x[0]*pi)', degree=1))
+
+    # Compute the adjoint gradient
+    adjoint_grad = y.gradient(m_prior=m2.vector().get_local(), source_term_prior=source_term2.vector().get_local())
+
+    # Compute the FD gradient
+    step = 1e-9   # finite diff step
+    FD_grad_m = optimize.approx_fprime(m2.vector().get_local(), lambda m2: y.logd(m2, source_term2.vector().get_local()), step)
+    FD_grad_source_term = optimize.approx_fprime(source_term2.vector().get_local(), lambda source_term2: y.logd(m2.vector().get_local(), source_term2), step)
+
+    # Assert gradients are close
+    assert np.allclose(adjoint_grad['m_prior'], FD_grad_m, rtol=2e-1),\
+        f"Adjoint gradient w.r.t. m: {adjoint_grad['m_prior']}, FD gradient w.r.t. m: {FD_grad_m}"
+    assert np.allclose(adjoint_grad['source_term_prior'], FD_grad_source_term, rtol=1e-1),\
+        f"Adjoint gradient w.r.t. source term: {adjoint_grad['source_term_prior']}, FD gradient w.r.t. source term: {FD_grad_source_term}"
+    assert np.linalg.norm(adjoint_grad['m_prior']-FD_grad_m)/ np.linalg.norm(FD_grad_m) < 1e-2
+    assert np.linalg.norm(adjoint_grad['source_term_prior']-FD_grad_source_term)/ np.linalg.norm(FD_grad_source_term) < 1e-2
+