Repository: CUQI-DTU/CUQIpy

Issue:
Multiple Inputs: Update forward model to accept multiple inputs 
**General description of task**

This is the second task outlined in the [design document](https://github.com/CUQI-DTU/CUQIpy-admin/blob/multiple_input_design/design/fwd_multiple_input.ipynb) for enabling multiple inputs. This task is to update the `model` module classes to enable multiple inputs for the forward map. It also address enabling gradient with respect to multiple inputs.




**Definition of Done (Feature/change)**
- [x] Classes of the `model` module are updated to support multiple inputs
- [x] (except PDE models to be differed to future issue) Classes of the `model` module are updated  to support gradient computation with respect to multiple inputs
- [x] (deferred to future issue due to the size of this PR) Raise error when model with multiple input is used within distributions (e.g. in the likelihood) since this is to be supported in upcoming tasks #649
- [x] (if applicable) The design file linked above is updated, this sub-feature use and functionality (it was updated).

**Definition of Done (Mandatory)**
- [x] Documentation added (docstrings on all public methods/classes)
- [x] Unit tests added/updated (and passing!) about 400 test cases added
- [x] Code reviewed and approved by 2 CUQI-maintainers
- [x] Online documentation is rendered correctly (esp. math) 

**Defintion of Done (Optional)**
- [x] (canceled, deferred for later tasks in the multiple input feature) HowTo/tutorial added

**additional notes**
Notes about the user interface and enabled features: please consult  [design document](https://github.com/CUQI-DTU/CUQIpy-admin/blob/multiple_input_design/design/fwd_multiple_input.ipynb) for the most recent added features
```python
## Forward map
# consider the forward model F we defined above
F(a_val, b_val, c_val) # should work (if passed as args, input should be in the same order 
                       # as what the domain geometry specifies, i.e. a, b, c)
                       # when possible, error is raised if the order is not correct
                       # this type of error can be detected, for example, when the 
                       # inputs a, b, and c have different parameter dimensions
F(a=a_val, b=b_val, c=c_val) # should work
F(b=b_val, c=c_val, a=a_val) # should work
# Future: support partial evaluation
F(a=a_val) # should work in the future, but not part of this task and returns
           # another forward model that is partially evaluated
#---------------------------------------------------
## Gradient
F.gradient(direction, a_val, b_val, c_val) # should work
```
~~# when possible also combination like below should work~~
~~F.gradient(direction, a=a_val) # should work~~
~~F.gradient(direction, a=a_val, c=c_val) # should work~~


**Notes about reviewing**
- I suggest looking at the [design document](https://github.com/CUQI-DTU/CUQIpy-admin/blob/multiple_input_design/design/fwd_multiple_input.ipynb) first, which I updated. The relevant part, is from the beginning up to the heading "The rest of the document is not yet updated". The content below this heading is not relevant to this review.
- The comparison file between Model class before and after the update helps identifying the methods that were added, [model_update.md](https://github.com/CUQI-DTU/CUQIpy-admin/blob/multiple_input_design/design/model_update.md).
- It might be better to review the tests before the model file because they showcase how to use the feature and what scenarios generate errors
- For the Model class, it is better to review the entire class (also referring to the diff file). Reviewing the diff file only can be confusing due to the so many changes
- For the remaining classes, LinearModel, PDEModel, etc., it makes sense to focus on the diff file only
- There are some remaining issues I will list them here: #648 #649

Hints from issue comments:
@amal-ghamdi: @jakobsj , DoD is ready for your review :)@amal-ghamdi: @jakobsj, DoD is ready for your review :)@jakobsj: @amal-ghamdi DoD approved!@amal-ghamdi: @jakobsj , Can you at some point take a look at the updated [design document](https://github.com/CUQI-DTU/CUQIpy-admin/blob/multiple_input_design/design/fwd_multiple_input.ipynb) up until the heading "The rest of the document is not yet updated". The content below this heading is not very relevant now and more or less have not been updated.

Many thanks!
Amal

Solution patch:
diff --git a/.github/ISSUE_TEMPLATE/task.md b/.github/ISSUE_TEMPLATE/task.md
new file mode 100644
index 000000000..eee701664
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/task.md
@@ -0,0 +1,16 @@
+**General description of task**
+
+<Insert description of change. Preferably with reason for change and a small (pseudo-)code example of the user interface>
+
+**Definition of Done (Feature/change)**
+- [ ] description1
+- [ ] description2
+
+**Definition of Done (Mandatory)**
+- [ ] Documentation added (docstrings on all public methods/classes)
+- [ ] Unit tests added/updated (and passing!)
+- [ ] Code reviewed and approved by 2 CUQI-maintainers
+- [ ] Online documentation is rendered correctly (esp. math) 
+
+**Defintion of Done (Optional)**
+- [ ] HowTo/tutorial added
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 4837bb0f0..8e9239bee 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -13,7 +13,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest] # Add windows-latest, macos-latest to test on multiple OS
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
diff --git a/cuqi/distribution/__init__.py b/cuqi/distribution/__init__.py
index 3a51f5472..4842e2cb4 100644
--- a/cuqi/distribution/__init__.py
+++ b/cuqi/distribution/__init__.py
@@ -14,6 +14,6 @@
 from ._normal import Normal
 from ._truncated_normal import TruncatedNormal
 from ._posterior import Posterior
-from ._uniform import Uniform
+from ._uniform import Uniform, UnboundedUniform
 from ._custom import UserDefinedDistribution, DistributionGallery
 from ._joint_distribution import JointDistribution, _StackedJointDistribution, MultipleLikelihoodPosterior
diff --git a/cuqi/distribution/_beta.py b/cuqi/distribution/_beta.py
index a57af1826..e4381d350 100644
--- a/cuqi/distribution/_beta.py
+++ b/cuqi/distribution/_beta.py
@@ -48,7 +48,7 @@ def logpdf(self, x):
 
         # Check bounds
         if np.any(x<=0) or np.any(x>=1) or np.any(self.alpha<=0) or np.any(self.beta<=0):
-            return -np.Inf
+            return -np.inf
 
         # Compute logpdf
         return np.sum(sps.beta.logpdf(x, a=self.alpha, b=self.beta))
diff --git a/cuqi/distribution/_cauchy.py b/cuqi/distribution/_cauchy.py
index bc8f166a7..80fa43102 100644
--- a/cuqi/distribution/_cauchy.py
+++ b/cuqi/distribution/_cauchy.py
@@ -75,14 +75,14 @@ def _is_out_of_bounds(self, x):
     def logpdf(self, x):
 
         if self._is_out_of_bounds(x):
-            return -np.Inf
+            return -np.inf
         
         return np.sum(-np.log(np.pi*self.scale*(1+((x-self.location)/self.scale)**2)))
     
     def cdf(self, x):
 
         if self._is_out_of_bounds(x):
-            return -np.Inf
+            return -np.inf
         
         return np.sum(sps.cauchy.cdf(x, loc=self.location, scale=self.scale))
     
diff --git a/cuqi/distribution/_truncated_normal.py b/cuqi/distribution/_truncated_normal.py
index ff9d253d7..51e97eb4d 100644
--- a/cuqi/distribution/_truncated_normal.py
+++ b/cuqi/distribution/_truncated_normal.py
@@ -37,7 +37,7 @@ class TruncatedNormal(Distribution):
         p = cuqi.distribution.TruncatedNormal(mean=0, std=1, low=-2, high=2)
         samples = p.sample(5000)
     """
-    def __init__(self, mean=None, std=None, low=-np.Inf, high=np.Inf, is_symmetric=False, **kwargs):
+    def __init__(self, mean=None, std=None, low=-np.inf, high=np.inf, is_symmetric=False, **kwargs):
         # Init from abstract distribution class
         super().__init__(is_symmetric=is_symmetric, **kwargs)  
 
@@ -97,7 +97,7 @@ def logpdf(self, x):
         # the unnormalized logpdf
         # check if x falls in the range between np.array a and b
         if np.any(x < self.low) or np.any(x > self.high):
-            return -np.Inf
+            return -np.inf
         else:
              return self._normal.logpdf(x)
 
@@ -107,7 +107,7 @@ def _gradient(self, x, *args, **kwargs):
         """
         # check if x falls in the range between np.array a and b
         if np.any(x < self.low) or np.any(x > self.high):
-            return np.NaN*np.ones_like(x)
+            return np.nan*np.ones_like(x)
         else:
             return self._normal.gradient(x, *args, **kwargs)
 
diff --git a/cuqi/distribution/_uniform.py b/cuqi/distribution/_uniform.py
index a61ae71e6..87d3054b1 100644
--- a/cuqi/distribution/_uniform.py
+++ b/cuqi/distribution/_uniform.py
@@ -1,5 +1,6 @@
 import numpy as np
 from cuqi.distribution import Distribution
+from cuqi.geometry import Geometry
 
 class Uniform(Distribution):
 
@@ -46,7 +47,7 @@ def gradient(self, x):
         Computes the gradient of logpdf at the given values of x.
         """
         if np.any(x < self.low) or np.any(x > self.high):
-            return np.NaN*np.ones_like(x)
+            return np.nan*np.ones_like(x)
         else:
             return np.zeros_like(x)
 
@@ -57,4 +58,37 @@ def _sample(self,N=1, rng=None):
         else:
             s = np.random.uniform(self.low, self.high, (N,self.dim)).T
 
-        return s
\ No newline at end of file
+        return s
+
+class UnboundedUniform(Distribution):
+    """
+    Unbounded uniform distribution. This is a special case of the
+    Uniform distribution, where the lower and upper bounds are set to
+    -inf and inf, respectively. This distribution is not normalizable,
+    and therefore cannot be sampled from. It is mainly used for
+    initializing non-informative priors.
+    Parameters
+    ----------
+    geometry : int or Geometry
+        The geometry of the distribution. If an integer is given, it is
+        interpreted as the dimension of the distribution. If a
+        Geometry object is given, its par_dim attribute is used.
+    """
+    def __init__(self, geometry, is_symmetric=True, **kwargs):
+        super().__init__(geometry=geometry, is_symmetric=is_symmetric, **kwargs) 
+
+    def logpdf(self, x):
+        """
+        Evaluate the logarithm of the unnormalized PDF at the given values of x.
+        """
+        # Always return 1.0 (the unnormalized log PDF)
+        return 1.0
+
+    def gradient(self, x):
+        """
+        Computes the gradient of logpdf at the given values of x.
+        """
+        return np.zeros_like(x)
+
+    def _sample(self, N=1, rng=None):
+        raise NotImplementedError("Cannot sample from UnboundedUniform distribution")
\ No newline at end of file
diff --git a/cuqi/experimental/geometry/_productgeometry.py b/cuqi/experimental/geometry/_productgeometry.py
index 1ff377327..3f4c0d490 100644
--- a/cuqi/experimental/geometry/_productgeometry.py
+++ b/cuqi/experimental/geometry/_productgeometry.py
@@ -172,10 +172,10 @@ def fun2vec(self, *funvals):
         return tuple(funvecs)
         
     
-    def __repr__(self) -> str:
+    def __repr__(self, pad="") -> str:
         """Representation of the product geometry."""
         string = "{}(".format(self.__class__.__name__) + "\n"
         for g in self.geometries:
-            string += "\t{}\n".format(g.__repr__())
-        string += ")"
+            string += pad + "    {}\n".format(g.__repr__())
+        string += pad + ")"
         return string
\ No newline at end of file
diff --git a/cuqi/experimental/mcmc/_conjugate.py b/cuqi/experimental/mcmc/_conjugate.py
index 3035cc3e7..e76c01e05 100644
--- a/cuqi/experimental/mcmc/_conjugate.py
+++ b/cuqi/experimental/mcmc/_conjugate.py
@@ -4,7 +4,7 @@
 from cuqi.experimental.mcmc import Sampler
 from cuqi.distribution import Posterior, Gaussian, Gamma, GMRF, ModifiedHalfNormal
 from cuqi.implicitprior import RegularizedGaussian, RegularizedGMRF, RegularizedUnboundedUniform
-from cuqi.utilities import get_non_default_args, count_nonzero, count_constant_components_1D, count_constant_components_2D
+from cuqi.utilities import get_non_default_args, count_nonzero, count_within_bounds, count_constant_components_1D, count_constant_components_2D, piecewise_linear_1D_DoF
 from cuqi.geometry import Continuous1D, Continuous2D, Image2D
 
 class Conjugate(Sampler):
@@ -17,8 +17,8 @@ class Conjugate(Sampler):
     - (GMRF, Gamma) where Gamma is defined on the precision parameter of the GMRF
     - (RegularizedGaussian, Gamma) with preset constraints only and Gamma is defined on the precision parameter of the RegularizedGaussian
     - (RegularizedGMRF, Gamma) with preset constraints only and Gamma is defined on the precision parameter of the RegularizedGMRF
-    - (RegularizedGaussian, ModifiedHalfNormal) with preset constraints and regularization only
-    - (RegularizedGMRF, ModifiedHalfNormal) with preset constraints and regularization only
+    - (RegularizedGaussian, ModifiedHalfNormal) with most of the preset constraints and regularization
+    - (RegularizedGMRF, ModifiedHalfNormal) with most of the preset constraints and regularization
 
     Currently the Gamma and ModifiedHalfNormal distribution must be univariate.
 
@@ -147,8 +147,8 @@ def validate_target(self):
         if self.target.prior.dim != 1:
             raise ValueError("RegularizedGaussian-Gamma conjugacy only works with univariate ModifiedHalfNormal prior")
 
-        if self.target.likelihood.distribution.preset["constraint"] not in ["nonnegativity"]:
-            raise ValueError("RegularizedGaussian-Gamma conjugacy only works with implicit regularized Gaussian likelihood with nonnegativity constraints")
+        # Raises error if preset is not supported
+        _compute_sparsity_level(self.target)
 
         key_value_pairs = _get_conjugate_parameter(self.target)
         if len(key_value_pairs) != 1:
@@ -166,7 +166,7 @@ def validate_target(self):
     def conjugate_distribution(self):
         # Extract variables
         b = self.target.likelihood.data                                 # mu
-        m = np.count_nonzero(b)                                         # n
+        m = _compute_sparsity_level(self.target)
         Ax = self.target.likelihood.distribution.mean                   # x_i
         L = self.target.likelihood.distribution(np.array([1])).sqrtprec # L
         alpha = self.target.prior.shape                                 # alpha
@@ -183,9 +183,9 @@ def validate_target(self):
         if self.target.prior.dim != 1:
             raise ValueError("RegularizedUnboundedUniform-Gamma conjugacy only works with univariate Gamma prior")
         
-        if self.target.likelihood.distribution.preset["regularization"] not in ["l1", "tv"]:
-            raise ValueError("RegularizedUnboundedUniform-Gamma conjugacy only works with implicit regularized Gaussian likelihood with l1 or tv regularization")
-        
+        # Raises error if preset is not supported
+        _compute_sparsity_level(self.target)
+
         key_value_pairs = _get_conjugate_parameter(self.target)
         if len(key_value_pairs) != 1:
             raise ValueError(f"Multiple references to conjugate parameter {self.target.prior.name} found in likelihood. Only one occurance is supported.")
@@ -219,8 +219,8 @@ def validate_target(self):
         if self.target.prior.dim != 1:
             raise ValueError("RegularizedGaussian-ModifiedHalfNormal conjugacy only works with univariate ModifiedHalfNormal prior")
         
-        if self.target.likelihood.distribution.preset["regularization"] not in ["l1", "tv"]:
-            raise ValueError("RegularizedGaussian-ModifiedHalfNormal conjugacy only works with implicit regularized Gaussian likelihood with l1 or tv regularization")
+        # Raises error if preset is not supported
+        _compute_sparsity_level(self.target)
 
         key_value_pairs = _get_conjugate_parameter(self.target)
         if len(key_value_pairs) != 2:
@@ -266,23 +266,74 @@ def conjugate_distribution(self):
     
 
 def _compute_sparsity_level(target):
-    """Computes the sparsity level in accordance with Section 4 from [2],"""
+    """Computes the sparsity level in accordance with Section 4 from [2],
+    this can be interpreted as the number of degrees of freedom, that is,
+    the number of components n minus the dimension the of the subdifferential of the regularized.
+    """
     x = target.likelihood.data
-    if target.likelihood.distribution.preset["constraint"] == "nonnegativity":
-        if target.likelihood.distribution.preset["regularization"] == "l1":
-            m = count_nonzero(x)
-        elif target.likelihood.distribution.preset["regularization"] == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
-            m = count_constant_components_1D(x, lower = 0.0)
-        elif target.likelihood.distribution.preset["regularization"] == "tv" and isinstance(target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
-            m = count_constant_components_2D(target.likelihood.distribution.geometry.par2fun(x), lower = 0.0)
-    else: # No constraints, only regularization
-        if target.likelihood.distribution.preset["regularization"] == "l1":
-            m = count_nonzero(x)
-        elif target.likelihood.distribution.preset["regularization"] == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
-            m = count_constant_components_1D(x)
-        elif target.likelihood.distribution.preset["regularization"] == "tv" and isinstance(target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
-            m = count_constant_components_2D(target.likelihood.distribution.geometry.par2fun(x))
-    return m
+
+    constraint = target.likelihood.distribution.preset["constraint"]
+    regularization = target.likelihood.distribution.preset["regularization"]
+
+    # There is no reference for some of these conjugacy rules
+    if constraint == "nonnegativity":
+        if regularization in [None, "l1"]:
+            # Number of non-zero components in x
+            return count_nonzero(x)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            # Number of non-zero constant components in x
+            return count_constant_components_1D(x, lower = 0.0)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
+            # Number of non-zero constant components in x
+            return count_constant_components_2D(target.likelihood.distribution.geometry.par2fun(x), lower = 0.0)
+    elif constraint == "box":
+        bounds = target.likelihood.distribution._box_bounds
+        if regularization is None:
+            # Number of components in x that are strictly between the lower and upper bound
+            return count_within_bounds(x, bounds[0], bounds[1])
+        elif regularization == "l1":
+            # Number of components in x that are strictly between the lower and upper bound and are not zero
+            return count_within_bounds(x, bounds[0], bounds[1], exception = 0.0)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            # Number of constant components in x between are strictly between the lower and upper bound
+            return count_constant_components_1D(x, lower = bounds[0], upper = bounds[1])
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
+            # Number of constant components in x between are strictly between the lower and upper bound
+            return count_constant_components_2D(target.likelihood.distribution.geometry.par2fun(x), lower = bounds[0], upper = bounds[1])
+    elif constraint in ["increasing", "decreasing"]:
+        if regularization is None:
+            # Number of constant components in x
+            return count_constant_components_1D(x)
+        elif regularization == "l1":
+            # Number of constant components in x that are not zero
+            return count_constant_components_1D(x, exception = 0.0)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            # Number of constant components in x
+            return count_constant_components_1D(x)
+        # Increasing and decreasing cannot be done in 2D
+    elif constraint in ["convex", "concave"]:
+        if regularization is None:
+            # Number of piecewise linear components in x
+            return piecewise_linear_1D_DoF(x)
+        elif regularization == "l1":
+            # Number of piecewise linear components in x that are not zero
+            return piecewise_linear_1D_DoF(x, exception_zero = True)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            # Number of piecewise linear components in x that are not flat
+            return piecewise_linear_1D_DoF(x, exception_flat = True)
+        # convex and concave has only been implemented in 1D
+    elif constraint == None: 
+        if regularization == "l1":
+            # Number of non-zero components in x
+            return count_nonzero(x)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            # Number of non-zero constant components in x
+            return count_constant_components_1D(x)
+        elif regularization == "tv" and isinstance(target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
+            # Number of non-zero constant components in x
+            return count_constant_components_2D(target.likelihood.distribution.geometry.par2fun(x))
+
+    raise ValueError("RegularizedGaussian preset constraint and regularization choice is currently not supported with conjugacy.")
 
 
 def _get_conjugate_parameter(target):
diff --git a/cuqi/experimental/mcmc/_cwmh.py b/cuqi/experimental/mcmc/_cwmh.py
index b05e0c9ec..84ee1e106 100644
--- a/cuqi/experimental/mcmc/_cwmh.py
+++ b/cuqi/experimental/mcmc/_cwmh.py
@@ -31,12 +31,9 @@ class CWMH(ProposalBasedSampler):
     initial_point : ndarray
         Initial parameters. *Optional*
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is
-        `callback(sample, sample_index)`, where `sample` is the current sample
-        and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
     kwargs : dict
         Additional keyword arguments to be passed to the base class 
diff --git a/cuqi/experimental/mcmc/_gibbs.py b/cuqi/experimental/mcmc/_gibbs.py
index e7673cdd8..ba23c585b 100644
--- a/cuqi/experimental/mcmc/_gibbs.py
+++ b/cuqi/experimental/mcmc/_gibbs.py
@@ -58,6 +58,10 @@ class HybridGibbs:
         will call its step method in each Gibbs step.
         Default is 1 for all variables.
 
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
+
     Example
     -------
     .. code-block:: python
@@ -103,7 +107,7 @@ class HybridGibbs:
             
     """
 
-    def __init__(self, target: JointDistribution, sampling_strategy: Dict[str, Sampler], num_sampling_steps: Dict[str, int] = None):
+    def __init__(self, target: JointDistribution, sampling_strategy: Dict[str, Sampler], num_sampling_steps: Dict[str, int] = None, callback=None):
 
         # Store target and allow conditioning to reduce to a single density
         self.target = target() # Create a copy of target distribution (to avoid modifying the original)
@@ -120,6 +124,9 @@ def __init__(self, target: JointDistribution, sampling_strategy: Dict[str, Sampl
         # Initialize sampler (after target is set)
         self._initialize()
 
+        # Set the callback function
+        self.callback = callback
+
     def _initialize(self):
         """ Initialize sampler """
 
@@ -158,13 +165,15 @@ def sample(self, Ns) -> 'HybridGibbs':
             The number of samples to draw.
 
         """
-
-        for _ in tqdm(range(Ns), "Sample: "):
+        for idx in tqdm(range(Ns), "Sample: "):
 
             self.step()
 
             self._store_samples()
 
+            # Call callback function if specified
+            self._call_callback(idx, Ns)
+
         return self
 
     def warmup(self, Nb, tune_freq=0.1) -> 'HybridGibbs':
@@ -192,6 +201,9 @@ def warmup(self, Nb, tune_freq=0.1) -> 'HybridGibbs':
                 
             self._store_samples()
 
+            # Call callback function if specified
+            self._call_callback(idx, Nb)
+
         return self
 
     def get_samples(self) -> Dict[str, Samples]:
@@ -263,6 +275,11 @@ def tune(self, skip_len, update_count):
             self.samplers[par_name].tune(skip_len=skip_len, update_count=update_count)
 
     # ------------ Private methods ------------
+    def _call_callback(self, sample_index, num_of_samples):
+        """ Calls the callback function. Assumes input is sampler, sample index, and total number of samples """
+        if self.callback is not None:
+            self.callback(self, sample_index, num_of_samples)
+
     def _initialize_samplers(self):
         """ Initialize samplers """
         for sampler in self.samplers.values():
diff --git a/cuqi/experimental/mcmc/_hmc.py b/cuqi/experimental/mcmc/_hmc.py
index 288805755..4c56bbe17 100644
--- a/cuqi/experimental/mcmc/_hmc.py
+++ b/cuqi/experimental/mcmc/_hmc.py
@@ -38,13 +38,9 @@ class NUTS(Sampler):
         opt_acc_rate should be in (0, 1), however, choosing a value that is very
         close to 1 or 0 might lead to poor performance of the sampler.
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is
-        `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of
-        the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
     Example
     -------
diff --git a/cuqi/experimental/mcmc/_langevin_algorithm.py b/cuqi/experimental/mcmc/_langevin_algorithm.py
index 5759a8548..021ac2bb2 100644
--- a/cuqi/experimental/mcmc/_langevin_algorithm.py
+++ b/cuqi/experimental/mcmc/_langevin_algorithm.py
@@ -32,11 +32,9 @@ class ULA(Sampler): # Refactor to Proposal-based sampler?
         be smaller than 1/L, where L is the Lipschitz of the gradient of the log
         target density, logd).
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
 
     Example
@@ -164,11 +162,9 @@ class MALA(ULA): # Refactor to Proposal-based sampler?
         be smaller than 1/L, where L is the Lipschitz of the gradient of the log
         target density, logd).
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
 
     Example
@@ -288,12 +284,9 @@ class MYULA(ULA):
     smoothing_strength : float
         This parameter controls the smoothing strength of MYULA.
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of
-        the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
     A Deblur example can be found in demos/howtos/myula.py
     # TODO: update demo once sampler merged
@@ -378,12 +371,9 @@ class PnPULA(MYULA):
         This parameter controls the smoothing strength of PnP-ULA.
 
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of
-        the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
     # TODO: update demo once sampler merged
     """
diff --git a/cuqi/experimental/mcmc/_laplace_approximation.py b/cuqi/experimental/mcmc/_laplace_approximation.py
index 649b613fb..b243609af 100644
--- a/cuqi/experimental/mcmc/_laplace_approximation.py
+++ b/cuqi/experimental/mcmc/_laplace_approximation.py
@@ -43,11 +43,9 @@ class UGLA(Sampler):
         sampling easier but results in a worse approximation. See details in Section 3.3 of the paper.
         If not provided, it defaults to 1e-5.
 
-    callback : callable, *Optional*
-        If set, this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
     """
     def __init__(self, target=None, initial_point=None, maxit=50, tol=1e-4, beta=1e-5, **kwargs):
 
diff --git a/cuqi/experimental/mcmc/_rto.py b/cuqi/experimental/mcmc/_rto.py
index 22d51d693..ad3593a5f 100644
--- a/cuqi/experimental/mcmc/_rto.py
+++ b/cuqi/experimental/mcmc/_rto.py
@@ -3,7 +3,7 @@
 from scipy.sparse.linalg import LinearOperator as scipyLinearOperator
 import numpy as np
 import cuqi
-from cuqi.solver import CGLS, FISTA, ADMM, ScipyLinearLSQ
+from cuqi.solver import CGLS, FISTA, ADMM, ScipyLinearLSQ, ScipyMinimizer
 from cuqi.experimental.mcmc import Sampler
 
 
@@ -36,11 +36,9 @@ class LinearRTO(Sampler):
     tol : float
         Tolerance of the inner CGLS solver. *Optional*.
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
         
     """
     def __init__(self, target=None, initial_point=None, maxit=10, tol=1e-6, **kwargs):
@@ -169,6 +167,7 @@ class RegularizedLinearRTO(LinearRTO):
     ADMM:  [2] Boyd et al. "Distributed optimization and statistical learning via the alternating direction method of multipliers."Foundations and Trends® in Machine learning, 2011.
            Used when prior.proximal is a list of penalty terms.
     ScipyLinearLSQ: Wrapper for Scipy's lsq_linear for the Trust Region Reflective algorithm. Optionally used when the constraint is either "nonnegativity" or "box".
+    ScipyMinimizer: Wrapper for Scipy's minimize. Optionally used when the constraint is either "nonnegativity" or "box".
 
     Parameters
     ------------
@@ -179,7 +178,7 @@ class RegularizedLinearRTO(LinearRTO):
         Initial point for the sampler. *Optional*.
 
     maxit : int
-        Maximum number of iterations of the FISTA/ADMM/ScipyLinearLSQ solver. *Optional*.
+        Maximum number of iterations of the FISTA/ADMM/ScipyLinearLSQ/ScipyMinimizer solver. *Optional*.
 
     inner_max_it : int
         Maximum number of iterations of the CGLS solver used within the ADMM solver. *Optional*.
@@ -193,7 +192,7 @@ class RegularizedLinearRTO(LinearRTO):
         See [2] or `cuqi.solver.ADMM`
 
     abstol : float
-        Absolute tolerance of the FISTA/ScipyLinearLSQ solver. *Optional*.
+        Absolute tolerance of the FISTA/ScipyLinearLSQ/ScipyMinimizer solver. *Optional*.
     
     inner_abstol : float
         Tolerance parameter for ScipyLinearLSQ's inner solve of the unbounded least-squares problem. *Optional*.
@@ -202,13 +201,11 @@ class RegularizedLinearRTO(LinearRTO):
         If True, FISTA is used as solver, otherwise ISTA is used. *Optional*.
     
     solver : string
-        If set to "ScipyLinearLSQ", solver is set to cuqi.solver.ScipyLinearLSQ, otherwise FISTA/ISTA or ADMM is used. Note "ScipyLinearLSQ" can only be used with `RegularizedGaussian` of `box` or `nonnegativity` constraint. *Optional*.
+        Options are "FISTA" (default for a single constraint or regularization), "ADMM" (default and the only option for multiple constraints or regularizations), "ScipyLinearLSQ" and "ScipyMinimizer". Note "ScipyLinearLSQ" and "ScipyMinimizer" can only be used with `RegularizedGaussian` of a single `box` or `nonnegativity` constraint. *Optional*.
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
         
     """
     def __init__(self, target=None, initial_point=None, maxit=100, inner_max_it=10, stepsize="automatic", penalty_parameter=10, abstol=1e-10, adaptive=True, solver=None, inner_abstol=None, **kwargs):
@@ -238,11 +235,11 @@ def solver(self):
 
     @solver.setter
     def solver(self, value):
-        if value == "ScipyLinearLSQ":
+        if value == "ScipyLinearLSQ" or value == "ScipyMinimizer":
             if (self.target.prior.preset["constraint"] == "nonnegativity" or self.target.prior.preset["constraint"] == "box"):
                 self._solver = value
             else:
-                raise ValueError("ScipyLinearLSQ only supports RegularizedGaussian with box or nonnegativity constraint.")
+                raise ValueError("ScipyLinearLSQ and ScipyMinimizer only support RegularizedGaussian with box or nonnegativity constraint.")
         else:
             self._solver = value
 
@@ -285,15 +282,22 @@ def step(self):
             sim = ADMM(self.M, y, self.proximal,
                         self.current_point, self.penalty_parameter, maxit = self.maxit, inner_max_it = self.inner_max_it, adaptive = self.adaptive)
         elif self.solver == "ScipyLinearLSQ":
-                A_op = sp.sparse.linalg.LinearOperator((sum([llh.dim for llh in self.likelihoods])+self.target.prior.dim, self.target.prior.dim),
-                                        matvec=lambda x: self.M(x, 1),
-                                        rmatvec=lambda x: self.M(x, 2)
-                                        )
-                sim = ScipyLinearLSQ(A_op, y, self.target.prior._box_bounds, 
-                                     max_iter = self.maxit,
-                                     lsmr_maxiter = self.inner_max_it, 
-                                     tol = self.abstol,
-                                     lsmr_tol = self.inner_abstol)
+            A_op = sp.sparse.linalg.LinearOperator((sum([llh.distribution.dim for llh in self.likelihoods])+self.target.prior.dim, self.target.prior.dim),
+                                    matvec=lambda x: self.M(x, 1),
+                                    rmatvec=lambda x: self.M(x, 2)
+                                    )
+            sim = ScipyLinearLSQ(A_op, y, self.target.prior._box_bounds, 
+                                    max_iter = self.maxit,
+                                    lsmr_maxiter = self.inner_max_it, 
+                                    tol = self.abstol,
+                                    lsmr_tol = self.inner_abstol)
+        elif self.solver == "ScipyMinimizer":
+            # Adapt bounds format, as scipy.minimize requires a bounds format 
+            # different than that in scipy.lsq_linear.
+            bounds = [(self.target.prior._box_bounds[0][i], self.target.prior._box_bounds[1][i]) for i in range(self.target.prior.dim)]
+            # Note that the objective function is defined as 0.5*||Mx-y||^2, 
+            # and the corresponding gradient (gradfunc) is given by M^T(Mx-y).
+            sim = ScipyMinimizer(lambda x: 0.5*np.sum((self.M(x, 1)-y)**2), self.current_point, gradfunc=lambda x: self.M(self.M(x, 1) - y, 2), bounds=bounds, tol=self.abstol, options={"maxiter": self.maxit})
         else:
             raise ValueError("Choice of solver not supported.")
 
diff --git a/cuqi/experimental/mcmc/_sampler.py b/cuqi/experimental/mcmc/_sampler.py
index b43b19ef7..90e11f2ee 100644
--- a/cuqi/experimental/mcmc/_sampler.py
+++ b/cuqi/experimental/mcmc/_sampler.py
@@ -59,9 +59,8 @@ def __init__(self, target:cuqi.density.Density=None, initial_point=None, callbac
             The initial point for the sampler. If not given, the sampler will choose an initial point.
 
         callback : callable, optional
-            A function that will be called after each sample is drawn. The function should take two arguments: the sample and the index of the sample.
-            The sample is a 1D numpy array and the index is an integer. The callback function is useful for monitoring the sampler during sampling.
-
+            A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+            The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
         """
 
         self.target = target
@@ -209,7 +208,6 @@ def sample(self, Ns, batch_size=0, sample_path='./CUQI_samples/') -> 'Sampler':
             The path to save the samples. If not specified, the samples are saved to the current working directory under a folder called 'CUQI_samples'.
 
         """
-
         self._ensure_initialized()
 
         # Initialize batch handler
@@ -235,7 +233,7 @@ def sample(self, Ns, batch_size=0, sample_path='./CUQI_samples/') -> 'Sampler':
                 batch_handler.add_sample(self.current_point)
 
             # Call callback function if specified            
-            self._call_callback(self.current_point, len(self._samples)-1)
+            self._call_callback(idx, Ns)
                 
         return self
     
@@ -276,7 +274,7 @@ def warmup(self, Nb, tune_freq=0.1) -> 'Sampler':
             pbar.set_postfix_str(f"acc rate: {np.mean(self._acc[-1-idx:]):.2%}")
 
             # Call callback function if specified
-            self._call_callback(self.current_point, len(self._samples)-1)
+            self._call_callback(idx, Nb)
 
         return self
     
@@ -367,10 +365,10 @@ def set_history(self, history: dict):
                 raise ValueError(f"Key {key} not recognized in history dictionary of sampler {self.__class__.__name__}.")
 
     # ------------ Private methods ------------
-    def _call_callback(self, sample, sample_index):
-        """ Calls the callback function. Assumes input is sample and sample index"""
+    def _call_callback(self, sample_index, num_of_samples):
+        """ Calls the callback function. Assumes input is sampler, sample index, and total number of samples """
         if self.callback is not None:
-            self.callback(sample, sample_index)
+            self.callback(self, sample_index, num_of_samples)
 
     def _validate_initialization(self):
         """ Validate the initialization of the sampler by checking all state and history keys are set. """
diff --git a/cuqi/geometry/_geometry.py b/cuqi/geometry/_geometry.py
index 3267ed89b..67be47dd0 100644
--- a/cuqi/geometry/_geometry.py
+++ b/cuqi/geometry/_geometry.py
@@ -225,7 +225,18 @@ def __eq__(self, obj):
         return self._all_values_equal(obj)
 
     def __repr__(self) -> str:
-        return "{}{}".format(self.__class__.__name__,self.par_shape)
+        if self.par_shape == self.fun_shape:
+            return "{}[{}]".format(self.__class__.__name__,
+                                   self.par_shape if len(self.par_shape) != 1 else self.par_shape[0])
+        return "{}[{}: {}]".format(
+            self.__class__.__name__,
+            self.par_shape if len(self.par_shape) != 1 else self.par_shape[0],
+            (
+                self.fun_shape
+                if (self.fun_shape is None or len(self.fun_shape) != 1)
+                else self.fun_shape[0]
+            ),
+        )
 
     def _all_values_equal(self, obj):
         """Returns true of all values of the object and self are equal"""
diff --git a/cuqi/implicitprior/_regularizedGaussian.py b/cuqi/implicitprior/_regularizedGaussian.py
index 53f86cc74..38a493aee 100644
--- a/cuqi/implicitprior/_regularizedGaussian.py
+++ b/cuqi/implicitprior/_regularizedGaussian.py
@@ -2,10 +2,12 @@
 from cuqi.distribution import Distribution, Gaussian
 from cuqi.solver import ProjectNonnegative, ProjectBox, ProximalL1
 from cuqi.geometry import Continuous1D, Continuous2D, Image2D
-from cuqi.operator import FirstOrderFiniteDifference, Operator
+from cuqi.operator import FirstOrderFiniteDifference, SecondOrderFiniteDifference, Operator
 
 import numpy as np
 import scipy.sparse as sparse
+import scipy.optimize as spoptimize
+
 from copy import copy
 
 
@@ -60,12 +62,18 @@ class RegularizedGaussian(Distribution):
         min_(z in C) 0.5||x-z||_2^2.
 
     constraint : string or None
-        Preset constraints that generate the corresponding proximal parameter. Can be set to "nonnegativity" and "box". Required for use in Gibbs.
-        For "box", the following additional parameters can be passed:
+        Preset constraints that generate the corresponding proximal parameter. Required for use in Gibbs. For any geometry the following can be chosen:
+        - "nonnegativity"
+        - "box", the following additional parameters can be passed:
             lower_bound : array_like or None
                 Lower bound of box, defaults to zero
             upper_bound : array_like
                 Upper bound of box, defaults to one
+        Additionally, for Continuous1D geometry the following options can be chosen:
+        - "increasing"
+        - "decreasing"
+        - "convex"
+        - "concave"
 
     regularization : string or None
         Preset regularization that generate the corresponding proximal parameter. Can be set to "l1" or 'tv'. Required for use in Gibbs in future update.
@@ -178,6 +186,34 @@ def _parse_preset_constraint_input(self, constraint, optional_regularization_par
                 self._proximal = lambda z, _: ProjectBox(z, _box_lower, _box_upper)
                 self._box_bounds = (np.ones(self.dim)*_box_lower, np.ones(self.dim)*_box_upper)
                 self._preset["constraint"] = "box"
+            elif c_lower == "increasing":
+                if not isinstance(self.geometry, Continuous1D):
+                    raise ValueError("Geometry not supported for " + c_lower)
+                if hasattr(spoptimize, 'isotonic_regression'):
+                    self._constraint_prox = lambda z, _: spoptimize.isotonic_regression(z, increasing=True).x
+                else:
+                    raise AttributeError(f"The function 'isotonic_regression' does not exist in scipy.optimize. Installed scipy version: {spoptimize.__version__}. You need to install a scipy >= 1.12.0")
+                self._preset["constraint"] = "increasing"
+            elif c_lower == "decreasing":
+                if not isinstance(self.geometry, Continuous1D):
+                    raise ValueError("Geometry not supported for " + c_lower)
+                if hasattr(spoptimize, 'isotonic_regression'):
+                    self._constraint_prox = lambda z, _: spoptimize.isotonic_regression(z, increasing=False).x
+                else:
+                    raise AttributeError(f"The function 'isotonic_regression' does not exist in scipy.optimize. Installed scipy version: {spoptimize.__version__}. You need to install a scipy >= 1.12.0")
+                self._preset["constraint"] = "decreasing"
+            elif c_lower == "convex":
+                if not isinstance(self.geometry, Continuous1D):
+                    raise ValueError("Geometry not supported for " + c_lower)
+                self._constraint_prox = lambda z, _: -ProjectNonnegative(-z)
+                self._constraint_oper = SecondOrderFiniteDifference(self.geometry.fun_shape, bc_type='neumann')
+                self._preset["constraint"] = "convex"
+            elif c_lower == "concave":
+                if not isinstance(self.geometry, Continuous1D):
+                    raise ValueError("Geometry not supported for " + c_lower)
+                self._constraint_prox = lambda z, _: ProjectNonnegative(z)
+                self._constraint_oper = SecondOrderFiniteDifference(self.geometry.fun_shape, bc_type='neumann')
+                self._preset["constraint"] = "concave"
             else:
                 raise ValueError("Constraint not supported.")
 
@@ -289,7 +325,7 @@ def _sample(self, N, rng=None):
   
     @staticmethod
     def constraint_options():
-        return ["nonnegativity", "box"]
+        return ["nonnegativity", "box", "increasing", "decreasing", "convex", "concave"]
 
     @staticmethod
     def regularization_options():
diff --git a/cuqi/model/_model.py b/cuqi/model/_model.py
index a811eaa6f..164938717 100644
--- a/cuqi/model/_model.py
+++ b/cuqi/model/_model.py
@@ -5,48 +5,55 @@
 from scipy.linalg import solve
 from cuqi.samples import Samples
 from cuqi.array import CUQIarray
-from cuqi.geometry import Geometry, _DefaultGeometry1D, _DefaultGeometry2D, _get_identity_geometries
+from cuqi.geometry import Geometry, _DefaultGeometry1D, _DefaultGeometry2D,\
+    _get_identity_geometries
 import cuqi
 import matplotlib.pyplot as plt
 from copy import copy
+from functools import partial
+from cuqi.utilities import force_ndarray
 
 class Model(object):
     """Generic model defined by a forward operator.
 
     Parameters
     -----------
-    forward : 2D ndarray or callable function.
-        Forward operator.
+    forward : callable function
+        Forward operator of the model. It takes one or more inputs and returns the model output.
 
-    range_geometry : integer or cuqi.geometry.Geometry
-        If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
+    range_geometry : integer, a 1D or 2D tuple of integers, cuqi.geometry.Geometry 
+        If integer or 1D tuple of integers is given, a cuqi.geometry._DefaultGeometry1D is created with dimension of the integer.
+        If 2D tuple of integers is given, a cuqi.geometry._DefaultGeometry2D is created with dimensions of the tuple.
+        If cuqi.geometry.Geometry object is given, it is used as the range geometry of the model.
 
-    domain_geometry : integer or cuqi.geometry.Geometry
-        If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
+    domain_geometry : integer, a 1D or 2D tuple of integers, cuqi.geometry.Geometry or a tuple with items of any of the listed types
+        If integer or 1D tuple of integers is given, a cuqi.geometry._DefaultGeometry1D is created with dimension of the integer.
+        If 2D tuple of integers is given (and the forward model has one input only), a cuqi.geometry._DefaultGeometry2D is created with dimensions of the tuple.
+        If cuqi.geometry.Geometry is given, it is used as the domain geometry.
+        If tuple of the above types is given, a cuqi.geometry._ProductGeometry is created based on the tuple entries. This is used for models with multiple inputs where each entry in the tuple represents the geometry of each input.
+
+    gradient : callable function, a tuple of callable functions or None, optional
+        The direction-Jacobian product of the forward model Jacobian with respect to the model input, evaluated at the model input. For example, if the forward model inputs are `x` and `y`, the gradient callable signature should be (`direction`, `x`, `y`), in that order, where `direction` is the direction by which the Jacobian matrix is multiplied and `x` and `y` are the parameters at which the Jacobian is computed.
+
+        If the gradient function is a single callable function, it returns a 1D ndarray if the model has only one input. If the model has multiple inputs, this gradient function should return a tuple of 1D ndarrays, each representing the gradient with respect to each input.
 
-    gradient : callable function, optional
-        The direction-Jacobian product of the forward operator Jacobian with 
-        respect to the forward operator input, evaluated at a point (`wrt`).
-        The signature of the gradient function should be (`direction`, `wrt`),
-        where `direction` is the direction by which the Jacobian matrix is
-        multiplied and `wrt` is the point at which the Jacobian is computed.
+        If the gradient function is a tuple of callable functions, each callable function should return a 1D ndarray representing the gradient with respect to each input. The order of the callable functions in the tuple should match the order of the model inputs.
 
-    jacobian : callable function, optional
-        The Jacobian of the forward operator with respect to the forward operator input,
-        evaluated at a point (`wrt`). The signature of the Jacobian function should be (`wrt`).
-        The Jacobian function should return a 2D ndarray of shape (range_dim, domain_dim).
-        The Jacobian function is used to specify the gradient function by computing the vector-Jacobian
-        product (VJP), here we refer to the vector in the VJP as the `direction` since it is the direction at 
-        which the gradient is computed.
-        automatically and thus the gradient function should not be specified when the Jacobian
-        function is specified.
+    jacobian : callable function, a tuple of callable functions or None, optional
+        The Jacobian of the forward model with respect to the forward model input, evaluated at the model input. For example, if the forward model inputs are `x` and `y`, the jacobian signature should be (`x`, `y`), in that order, where `x` and `y` are the parameters at which the Jacobian is computed.
+
+        If the Jacobian function is a single callable function, it should return a 2D ndarray of shape (range_dim, domain_dim) if the model has only one input. If the model has multiple inputs, this Jacobian function should return a tuple of 2D ndarrays, each representing the Jacobian with respect to each input. 
+        
+        If the Jacobian function is a tuple of callable functions, each callable function should return a 2D ndarray representing the Jacobian with respect to each input. The order of the callable functions in the tuple should match the order of the model inputs.
+    
+        The Jacobian function is used to specify the gradient function by computing the vector-Jacobian product (VJP), here we refer to the vector in the VJP as the `direction` since it is the direction at which the gradient is computed. Either the gradient or the Jacobian can be specified, but not both.
 
 
     :ivar range_geometry: The geometry representing the range.
     :ivar domain_geometry: The geometry representing the domain.
 
-    Example
-    -------
+    Example 1
+    ----------
 
     Consider a forward model :math:`F: \mathbb{R}^2 \\rightarrow \mathbb{R}` defined by the following forward operator:
 
@@ -75,6 +82,9 @@ def jacobian(x): # Can use "x" or "wrt" as the input argument name
 
         model = Model(forward, range_geometry=1, domain_geometry=2, jacobian=jacobian)
 
+        print(model(np.array([1, 1])))
+        print(model.gradient(np.array([1]), np.array([1, 1])))
+
     Alternatively, the gradient information in the forward model can be defined by direction-Jacobian product using the gradient keyword argument.
 
     This may be more efficient if forming the Jacobian matrix is expensive.
@@ -87,64 +97,158 @@ def jacobian(x): # Can use "x" or "wrt" as the input argument name
         def forward(x):
             return 10*x[1] - 10*x[0]**3 + 5*x[0]**2 + 6*x[0]
 
-        def gradient(direction, wrt):
-            # Direction-Jacobian product direction@jacobian(wrt)
-            return direction@np.array([[-30*wrt[0]**2 + 10*wrt[0] + 6, 10]])
+        def gradient(direction, x):
+            # Direction-Jacobian product direction@jacobian(x)
+            return direction@np.array([[-30*x[0]**2 + 10*x[0] + 6, 10]])
 
         model = Model(forward, range_geometry=1, domain_geometry=2, gradient=gradient)
 
+        print(model(np.array([1, 1])))
+        print(model.gradient(np.array([1]), np.array([1, 1])))
+
+    Example 2
+    ----------
+    Alternatively, the example above can be defined as a model with multiple inputs: :math:`x` and :math:`y`:
+
+    .. code-block:: python
+
+        import numpy as np
+        from cuqi.model import Model
+        from cuqi.geometry import Discrete
+    
+        def forward(x, y):
+            return 10 * y - 10 * x**3 + 5 * x**2 + 6 * x
+    
+        def jacobian(x, y):
+            return (np.array([[-30 * x**2 + 10 * x + 6]]), np.array([[10]]))
+    
+        model = Model(
+            forward,
+            range_geometry=1,
+            domain_geometry=(Discrete(1), Discrete(1)),
+            jacobian=jacobian,
+        )
+    
+        print(model(1, 1))
+        print(model.gradient(np.array([1]), 1, 1))
     """
     def __init__(self, forward, range_geometry, domain_geometry, gradient=None, jacobian=None):
 
-        #Check if input is callable
+        # Check if input is callable
         if callable(forward) is not True:
             raise TypeError("Forward needs to be callable function.")
-        
+
+        # Store forward func
+        self._forward_func = forward
+        self._stored_non_default_args = None
+
+        # Store range_geometry
+        self.range_geometry = range_geometry
+
+        # Store domain_geometry
+        self.domain_geometry = domain_geometry
+
+        # Additional checks for the forward operator
+        self._check_domain_geometry_consistent_with_forward()
+
         # Check if only one of gradient and jacobian is given
         if (gradient is not None) and (jacobian is not None):
             raise TypeError("Only one of gradient and jacobian should be specified")
-        
-        #Check if input is callable
-        if (gradient is not None) and (callable(gradient) is not True):
-            raise TypeError("Gradient needs to be callable function.")
-        
-        if (jacobian is not None) and (callable(jacobian) is not True):
-            raise TypeError("Jacobian needs to be callable function.")
-        
-        # Use jacobian function to specify gradient function (vector-Jacobian product)
+
+        # Check correct gradient form (check type, signature, etc.)
+        self._check_correct_gradient_jacobian_form(gradient, "gradient")
+
+        # Check correct jacobian form (check type, signature, etc.)
+        self._check_correct_gradient_jacobian_form(jacobian, "jacobian")
+
+        # If jacobian is provided, use it to specify gradient function
+        # (vector-Jacobian product)
         if jacobian is not None:
-            gradient = lambda direction, wrt: direction@jacobian(wrt)
- 
-        #Store forward func
-        self._forward_func = forward
+            gradient = self._use_jacobian_to_specify_gradient(jacobian)
+
         self._gradient_func = gradient
-         
-        #Store range_geometry
-        if isinstance(range_geometry, tuple) and len(range_geometry) == 2:
-            self.range_geometry = _DefaultGeometry2D(range_geometry)
-        elif isinstance(range_geometry, int):
-            self.range_geometry = _DefaultGeometry1D(grid=range_geometry)
-        elif isinstance(range_geometry, Geometry):
-            self.range_geometry = range_geometry
-        elif range_geometry is None:
-            raise AttributeError("The parameter 'range_geometry' is not specified by the user and it connot be inferred from the attribute 'forward'.")
-        else:
-            raise TypeError("The parameter 'range_geometry' should be of type 'int', 2 dimensional 'tuple' or 'cuqi.geometry.Geometry'.")
-
-        #Store domain_geometry
-        if isinstance(domain_geometry, tuple) and len(domain_geometry) == 2:
-            self.domain_geometry = _DefaultGeometry2D(domain_geometry)
-        elif isinstance(domain_geometry, int):
-            self.domain_geometry = _DefaultGeometry1D(grid=domain_geometry)
-        elif isinstance(domain_geometry, Geometry):
-            self.domain_geometry = domain_geometry
-        elif domain_geometry is None:
-            raise AttributeError("The parameter 'domain_geometry' is not specified by the user and it connot be inferred from the attribute 'forward'.")
+
+        # Set gradient output stacked flag to False
+        self._gradient_output_stacked = False
+
+    @property
+    def _non_default_args(self):
+        if self._stored_non_default_args is None:
+            # Store non_default_args of the forward operator for faster caching
+            # when checking for those arguments.
+            self._stored_non_default_args =\
+                cuqi.utilities.get_non_default_args(self._forward_func)
+        return self._stored_non_default_args
+ 
+    @property
+    def number_of_inputs(self):
+        """ The number of inputs of the model. """
+        return len(self._non_default_args)
+
+    @property
+    def range_geometry(self):
+        """ The geometry representing the range of the model. """
+        return self._range_geometry
+
+    @range_geometry.setter
+    def range_geometry(self, value):
+        """ Update the range geometry of the model. """
+        if isinstance(value, Geometry):
+            self._range_geometry = value
+        elif isinstance(value, int):
+            self._range_geometry = self._create_default_geometry(value)
+        elif isinstance(value, tuple):
+            self._range_geometry = self._create_default_geometry(value)
+        elif value is None:
+            raise AttributeError(
+                "The parameter 'range_geometry' is not specified by the user and it cannot be inferred from the attribute 'forward'."
+            )
         else:
-            raise TypeError("The parameter 'domain_geometry' should be of type 'int', 2 dimensional 'tuple' or 'cuqi.geometry.Geometry'.")
+            raise TypeError(
+                " The allowed types for 'range_geometry' are: 'cuqi.geometry.Geometry', int, 1D tuple of int, or 2D tuple of int."
+            )
 
-        # Store non_default_args of the forward operator for faster caching when checking for those arguments.
-        self._non_default_args = cuqi.utilities.get_non_default_args(self._forward_func)
+    @property
+    def domain_geometry(self):
+        """ The geometry representing the domain of the model. """
+        return self._domain_geometry
+
+    @domain_geometry.setter
+    def domain_geometry(self, value):
+        """ Update the domain geometry of the model. """
+
+        if isinstance(value, Geometry):
+            self._domain_geometry = value
+        elif isinstance(value, int):
+            self._domain_geometry = self._create_default_geometry(value)
+        elif isinstance(value, tuple) and self.number_of_inputs == 1:
+            self._domain_geometry = self._create_default_geometry(value)
+        elif isinstance(value, tuple) and self.number_of_inputs > 1:
+            geometries = [item if isinstance(item, Geometry) else self._create_default_geometry(item) for item in value]
+            self._domain_geometry = cuqi.experimental.geometry._ProductGeometry(*geometries)
+        elif value is None:
+            raise AttributeError(
+                "The parameter 'domain_geometry' is not specified by the user and it cannot be inferred from the attribute 'forward'."
+            )
+        else:
+            raise TypeError(
+                "For forward model with 1 input, the allowed types for 'domain_geometry' are: 'cuqi.geometry.Geometry', int, 1D tuple of int, or 2D tuple of int. For forward model with multiple inputs, the 'domain_geometry' should be a tuple with items of any of the above types."
+            )
+
+    def _create_default_geometry(self, value):
+        """Private function that creates default geometries for the model."""
+        if isinstance(value, tuple) and len(value) == 1:
+            value = value[0]
+        if isinstance(value, Geometry):
+            return value
+        if isinstance(value, int):
+            return _DefaultGeometry1D(grid=value)
+        elif isinstance(value, tuple) and len(value) == 2:
+            return _DefaultGeometry2D(im_shape=value)
+        else:
+            raise ValueError(
+                "Default geometry creation can be specified by an integer or a 2D tuple of integers."
+            )
 
     @property
     def domain_dim(self):
@@ -160,341 +264,842 @@ def range_dim(self):
         """
         return self.range_geometry.par_dim
 
-    def _2fun(self, x, geometry, is_par):
-        """ Converts `x` to function values (if needed) using the appropriate 
-        geometry. For example, `x` can be the model input which need to be
-        converted to function value before being passed to 
-        :class:`~cuqi.model.Model` operators (e.g. _forward_func, _adjoint_func,
-        _gradient_func).
+    def _check_domain_geometry_consistent_with_forward(self):
+        """Private function that checks if the domain geometry of the model is
+        consistent with the forward operator."""
+        if (
+            not isinstance(
+                self.domain_geometry, cuqi.experimental.geometry._ProductGeometry
+            )
+            and self.number_of_inputs > 1
+        ):
+            raise ValueError(
+                "The forward operator input is specified by more than one argument. This is only supported for domain geometry of type tuple with items of type: cuqi.geometry.Geometry object, int, or 2D tuple of int."
+            )
+
+    def _check_correct_gradient_jacobian_form(self, func, func_type):
+        """Private function that checks if the gradient/jacobian parameter is
+        in the correct form. That is, check if the gradient/jacobian has the
+        correct type, signature, etc."""
+
+        if func is None:
+            return
+
+        # gradient/jacobian should be callable (for single input and multiple input case)
+        # or a tuple of callables (for multiple inputs case)
+        if isinstance(func, tuple):
+            # tuple length should be same as the number of inputs
+            if len(func) != self.number_of_inputs:
+                raise ValueError(
+                    f"The "
+                    + func_type.lower()
+                    + f" tuple length should be {self.number_of_inputs} for model with inputs {self._non_default_args}"
+                )
+            # tuple items should be callables or None
+            if not all([callable(func_i) or func_i is None for func_i in func]):
+                raise TypeError(
+                    func_type.capitalize()
+                    + " tuple should contain callable functions or None."
+                )
+
+        elif callable(func):
+            # temporarily convert gradient/jacobian to tuple for checking only
+            func = (func,)
+
+        else:
+            raise TypeError(
+                "Gradient needs to be callable function or tuple of callable functions."
+            )
+
+        expected_func_non_default_args = self._non_default_args
+        if func_type.lower() == "gradient":
+            # prepend 'direction' to the expected gradient non default args
+            expected_func_non_default_args = [
+                "direction"
+            ] + expected_func_non_default_args
+
+        for func_i in func:
+            # make sure the signature of the gradient/jacobian function is correct
+            # that is, the same as the expected_func_non_default_args
+            if func_i is not None:
+                func_non_default_args = cuqi.utilities.get_non_default_args(func_i)
+
+                if list(func_non_default_args) != list(expected_func_non_default_args):
+                    raise ValueError(
+                        func_type.capitalize()
+                        + f" function signature should be {expected_func_non_default_args}"
+                    )
+
+    def _use_jacobian_to_specify_gradient(self, jacobian):
+        """Private function that uses the jacobian function to specify the
+        gradient function."""
+        # if jacobian is a single function and model has multiple inputs
+        if callable(jacobian) and self.number_of_inputs > 1:
+            gradient = self._create_gradient_lambda_function_from_jacobian_with_correct_signature(
+                jacobian, form='one_callable_multiple_inputs'
+            )
+        # Elif jacobian is a single function and model has only one input
+        elif callable(jacobian):
+            gradient = self._create_gradient_lambda_function_from_jacobian_with_correct_signature(
+                jacobian, form='one_callable_one_input'
+            )
+        # Else, jacobian is a tuple of jacobian functions
+        else:
+            gradient = []
+            for jac in jacobian:
+                if jac is not None:
+                    gradient.append(
+                        self._create_gradient_lambda_function_from_jacobian_with_correct_signature(
+                            jac, form='tuple_of_callables'
+                        )
+                    )
+                else:
+                    gradient.append(None)
+        return tuple(gradient) if isinstance(gradient, list) else gradient
+
+    def _create_gradient_lambda_function_from_jacobian_with_correct_signature(
+        self, jacobian, form
+    ):
+        """Private function that creates gradient lambda function from the
+        jacobian function, with the correct signature (based on the model
+        non_default_args).
+        """
+        # create the string representation of the lambda function
+        # for different forms of jacobian
+        if form=='one_callable_multiple_inputs':
+            grad_fun_str = (
+                "lambda direction, "
+                + ", ".join(self._non_default_args)
+                + ", jacobian: tuple([direction@jacobian("
+                + ", ".join(self._non_default_args)
+                + ")[i] for i in range("+str(self.number_of_inputs)+")])"
+            )
+        elif form=='tuple_of_callables' or form=='one_callable_one_input':
+            grad_fun_str = (
+                "lambda direction, "
+                + ", ".join(self._non_default_args)
+                + ", jacobian: direction@jacobian("
+                + ", ".join(self._non_default_args)
+                + ")"
+            )
+        else:
+            raise ValueError("form should be either 'one_callable' or 'tuple_of_callables'.")
+
+        # create the lambda function from the string
+        grad_func = eval(grad_fun_str)
+
+        # create partial function from the lambda function with jacobian as a
+        # fixed argument
+        grad_func = partial(grad_func, jacobian=jacobian)
+
+        return grad_func
+
+    def _2fun(self, geometry=None, is_par=True, **kwargs):
+        """ Converts `kwargs` to function values (if needed) using the geometry. For example, `kwargs` can be the model input which need to be converted to function value before being passed to :class:`~cuqi.model.Model` operators (e.g. _forward_func, _adjoint_func, _gradient_func).
 
         Parameters
         ----------
-        x : ndarray or cuqi.array.CUQIarray
-            The value to be converted.
-
         geometry : cuqi.geometry.Geometry
-            The geometry representing `x`.
+            The geometry representing the values in `kwargs`.
 
-        is_par : bool
-            If True, `x` is assumed to be parameters.
-            If False, `x` is assumed to be function values.
+        is_par : bool or a tuple of bools
+            If `is_par` is True, the values in `kwargs` are assumed to be parameters.
+            If `is_par` is False, the values in `kwargs` are assumed to be function values.
+            If `is_par` is a tuple of bools, the values in `kwargs` are assumed to be parameters or function values based on the corresponding boolean value in the tuple.
+        
+        **kwargs : keyword arguments to be converted to function values.
 
         Returns
         -------
-        ndarray or cuqi.array.CUQIarray
-            `x` represented as a function.
+        dict of the converted values
         """
-        # Convert to function representation
-        # if x is CUQIarray and geometry are consistent, we obtain funvals
-        # directly
-        if isinstance(x, CUQIarray) and  x.geometry == geometry:
-            x = x.funvals
-        # Otherwise we use the geometry par2fun method
-        elif is_par:
-            x = geometry.par2fun(x)
-
-        return x
-
-    def _2par(self, val, geometry, to_CUQIarray=False, is_par=False):
-        """ Converts val, normally output of :class:~`cuqi.model.Model` 
-        operators (e.g. _forward_func, _adjoint_func, _gradient_func), to
-        parameters using the appropriate geometry.
+        # Check kwargs and geometry are consistent and set up geometries list and
+        # is_par tuple
+        geometries, is_par = self._helper_pre_conversion_checks_and_processing(geometry, is_par, **kwargs)
+
+        # Convert to function values
+        for i, (k, v) in enumerate(kwargs.items()):
+            # Use CUQIarray funvals if geometry is consistent
+            if isinstance(v, CUQIarray) and v.geometry == geometries[i]:
+                kwargs[k] = v.funvals
+            # Else, if we still need to convert to function value (is_par[i] is True) 
+            # we use the geometry par2fun method
+            elif is_par[i] and v is not None:
+                kwargs[k] = geometries[i].par2fun(v)
+            else:
+                # No need to convert
+                pass
+
+        return kwargs
+
+    def _helper_pre_conversion_checks_and_processing(self, geometry=None, is_par=True, **kwargs):
+        """ Helper function that checks if kwargs and geometry are consistent
+        and sets up geometries list and is_par tuple.
+        """
+        # If len of kwargs is larger than 1, the geometry needs to be of type
+        # _ProductGeometry
+        if (
+            not isinstance(geometry, cuqi.experimental.geometry._ProductGeometry)
+            and len(kwargs) > 1
+        ):
+            raise ValueError(
+                "The input is specified by more than one argument. This is only "
+                + "supported for domain geometry of type "
+                + f"{cuqi.experimental.geometry._ProductGeometry.__name__}."
+            )
+
+        # If is_par is bool, make it a tuple of bools of the same length as
+        # kwargs
+        is_par = (is_par,) * len(kwargs) if isinstance(is_par, bool) else is_par
+
+        # Set up geometries list
+        geometries = (
+            geometry.geometries
+            if isinstance(geometry, cuqi.experimental.geometry._ProductGeometry)
+            else [geometry]
+        )
+
+        return geometries, is_par
+
+    def _2par(self, geometry=None, to_CUQIarray=False, is_par=False, **kwargs):    
+        """ Converts `kwargs` to parameters using the geometry. For example, `kwargs` can be the output of :class:`~cuqi.model.Model` operators (e.g. _forward_func, _adjoint_func, _gradient_func) which need to be converted to parameters before being returned.
 
         Parameters
         ----------
-        val : ndarray or cuqi.array.CUQIarray
-            The value to be converted to parameters.
-
         geometry : cuqi.geometry.Geometry
-            The geometry representing the argument `val`.
+            The geometry representing the values in `kwargs`.
 
-        to_CUQIarray : bool
-            If True, the returned value is wrapped as a cuqi.array.CUQIarray.
+        to_CUQIarray : bool or a tuple of bools
+            If `to_CUQIarray` is True, the values in `kwargs` will be wrapped in `CUQIarray`.
+            If `to_CUQIarray` is False, the values in `kwargs` will not be wrapped in `CUQIarray`.
+            If `to_CUQIarray` is a tuple of bools, the values in `kwargs` will be wrapped in `CUQIarray` or not based on the corresponding boolean value in the tuple.
         
-        is_par : bool
-            If True, `val` is assumed to be of parameter representation and
-            hence no conversion to parameters is performed.
+        is_par : bool or a tuple of bools
+            If `is_par` is True, the values in `kwargs` are assumed to be parameters.
+            If `is_par` is False, the values in `kwargs` are assumed to be function values.
+            If `is_par` is a tuple of bools, the values in `kwargs` are assumed to be parameters or function values based on the corresponding boolean value in the tuple.
 
         Returns
         -------
-        ndarray or cuqi.array.CUQIarray
-            The value `val` represented as parameters.
+        dict of the converted values
         """
+        # Check kwargs and geometry are consistent and set up geometries list and
+        # is_par tuple
+        geometries, is_par = self._helper_pre_conversion_checks_and_processing(geometry, is_par, **kwargs)
+
+        # if to_CUQIarray is bool, make it a tuple of bools of the same length
+        # as kwargs
+        to_CUQIarray = (to_CUQIarray,) * len(kwargs) if isinstance(to_CUQIarray, bool) else to_CUQIarray
+
         # Convert to parameters
-        # if val is CUQIarray and geometry are consistent, we obtain parameters
-        # directly
-        if isinstance(val, CUQIarray) and val.geometry == geometry:
-            val = val.parameters
-        # Otherwise we use the geometry fun2par method
-        elif not is_par:
-            val = geometry.fun2par(val)
-
-        # Wrap val in CUQIarray if requested
-        if to_CUQIarray:
-            val = CUQIarray(val, is_par=True, geometry=geometry)
-
-        # Return val
-        return val
-        
+        for i , (k, v) in enumerate(kwargs.items()):
+            # Use CUQIarray parameters if geometry is consistent
+            if isinstance(v, CUQIarray) and v.geometry == geometries[i]:
+                v = v.parameters
+            # Else, if we still need to convert to parameter value (is_par[i] is False) 
+            # we use the geometry fun2par method
+            elif not is_par[i] and v is not None:
+                v = geometries[i].fun2par(v)
+            else:
+                # No need to convert
+                pass
+
+            # Wrap the value v in CUQIarray if requested
+            if to_CUQIarray[i] and v is not None:
+                v = CUQIarray(v, is_par=True, geometry=geometries[i])
+
+            kwargs[k] = v
 
-    def _apply_func(self, func, func_range_geometry, func_domain_geometry, x, is_par, **kwargs):
-        """ Private function that applies the given function `func` to the input value `x`. It converts the input to function values (if needed) using the given `func_domain_geometry` and converts the output function values to parameters using the given `func_range_geometry`. It additionally handles the case of applying the function `func` to the cuqi.samples.Samples object.
+        return kwargs
 
-        kwargs are keyword arguments passed to the functions `func`.
+    def _apply_func(self, func=None, fwd=True, is_par=True, **kwargs):
+        """ Private function that applies the given function `func` to the input `kwargs`. It converts the input to function values (if needed) and converts the output to parameter values. It additionally handles the case of applying the function `func` to cuqi.samples.Samples objects.
         
         Parameters
         ----------
         func: function handler 
             The function to be applied.
 
-        func_range_geometry : cuqi.geometry.Geometry
-            The geometry representing the function `func` range.
-
-        func_domain_geometry : cuqi.geometry.Geometry
-            The geometry representing the function `func` domain.
-
-        x : ndarray or cuqi.array.CUQIarray
-            The input value to the operator.
+        fwd : bool
+            Flag indicating the direction of the operator to determine the range and domain geometries of the function.
+            If True the function is a forward operator.
+            If False the function is an adjoint operator.
 
-        is_par : bool
-            If True the input is assumed to be parameters.
-            If False the input is assumed to be function values.
+        is_par : bool or list of bool
+            If True, the inputs in `kwargs` are assumed to be parameters.
+            If False, the input in `kwargs` are assumed to be function values.
+            If `is_par` is a list of bools, the inputs are assumed to be parameters or function values based on the corresponding boolean value in the list.
 
         Returns
         -------
-        ndarray or cuqi.array.CUQIarray
-            The output of the function `func` converted to parameters.
+        ndarray or cuqi.array.CUQIarray or cuqi.samples.Samples object
+            The output of the function.
         """ 
+        # Specify the range and domain geometries of the function
+        # If forward operator, range geometry is the model range geometry and
+        # domain geometry is the model domain geometry
+        if fwd:
+            func_range_geometry = self.range_geometry
+            func_domain_geometry = self.domain_geometry
+        # If adjoint operator, range geometry is the model domain geometry and
+        # domain geometry is the model range geometry
+        else:
+            func_range_geometry = self.domain_geometry
+            func_domain_geometry = self.range_geometry
+
         # If input x is Samples we apply func for each sample
         # TODO: Check if this can be done all-at-once for computational speed-up
-        if isinstance(x,Samples):
-            out = np.zeros((func_range_geometry.par_dim, x.Ns))
-            # Recursively apply func to each sample
-            for idx, item in enumerate(x):
-                out[:,idx] = self._apply_func(func,
-                                              func_range_geometry,
-                                              func_domain_geometry,
-                                              item, is_par=True,
-                                              **kwargs)
-            return Samples(out, geometry=func_range_geometry)
-        
-        # store if input x is CUQIarray
-        is_CUQIarray = type(x) is CUQIarray
+        if any(isinstance(x, Samples) for x in kwargs.values()):
+            return self._handle_case_when_model_input_is_samples(func, fwd, **kwargs)
 
-        x = self._2fun(x, func_domain_geometry, is_par=is_par)
-        out = func(x, **kwargs)
+        # store if any input x is CUQIarray
+        is_CUQIarray = any(isinstance(x, CUQIarray) for x in kwargs.values())
 
-        # Return output as parameters 
-        # (and wrapped in CUQIarray if input was CUQIarray)
-        return self._2par(out, func_range_geometry, 
-                                    to_CUQIarray=is_CUQIarray)
+        # Convert input to function values
+        kwargs = self._2fun(geometry=func_domain_geometry, is_par=is_par, **kwargs)
 
-    def _parse_args_add_to_kwargs(self, *args, **kwargs):
-        """ Private function that parses the input arguments of the model and adds them as keyword arguments matching the non default arguments of the forward function. """
+        # Apply the function
+        out = func(**kwargs)
 
-        if len(args) > 0:
+        # Return output as parameters
+        # (wrapped in CUQIarray if any input was CUQIarray)
+        return self._2par(
+            geometry=func_range_geometry, to_CUQIarray=is_CUQIarray, **{"out": out}
+        )["out"]
 
+    def _handle_case_when_model_input_is_samples(self, func=None, fwd=True, **kwargs):
+        """Private function that calls apply_func for samples in the
+        Samples object(s).
+        """
+        # All kwargs should be Samples objects
+        if not all(isinstance(x, Samples) for x in kwargs.values()):
+            raise TypeError(
+                "If applying the function to Samples, all inputs should be Samples."
+            )
+
+        # All Samples objects should have the same number of samples
+        Ns = list(kwargs.values())[0].Ns
+        if not all(x.Ns == Ns for x in kwargs.values()):
+            raise ValueError(
+                "If applying the function to Samples, all inputs should have the same number of samples."
+            )
+
+        # Specify the range dimension of the function
+        range_dim = self.range_dim if fwd else self.domain_dim
+
+        # Create empty array to store the output
+        out = np.zeros((range_dim, Ns))
+
+        # Recursively apply func to each sample
+        for i in range(Ns):
+            kwargs_i = {
+                k: CUQIarray(v.samples[..., i], is_par=v.is_par, geometry=v.geometry)
+                for k, v in kwargs.items()
+            }
+            out[:, i] = self._apply_func(func=func, fwd=fwd, **kwargs_i)
+        # Specify the range geometries of the function
+        func_range_geometry = self.range_geometry if fwd else self.domain_geometry
+        return Samples(out, geometry=func_range_geometry)
+
+    def _parse_args_add_to_kwargs(
+        self, *args, is_par=True, non_default_args=None, map_name="model", **kwargs
+    ):
+        """ Private function that parses the input arguments and adds them as
+        keyword arguments matching (the order of) the non default arguments of
+        the forward function or other specified non_default_args list.
+        """
+        # If non_default_args is not specified, use the non_default_args of the
+        # model
+        if non_default_args is None:
+            non_default_args = self._non_default_args
+
+        # If any args are given, add them to kwargs
+        if len(args) > 0:
             if len(kwargs) > 0:
-                raise ValueError("The model input is specified both as positional and keyword arguments. This is not supported.")
-                
-            if len(args) != len(self._non_default_args):
-                raise ValueError("The number of positional arguments does not match the number of non-default arguments of the model.")
-            
+                raise ValueError(
+                    "The "
+                    + map_name.lower()
+                    + " input is specified both as positional and keyword arguments. This is not supported."
+                )
+
+            appending_error_message = ""
+            # Check if the input is for multiple input case and is stacked,
+            # then split it
+            if len(args)==1 and len(non_default_args)>1:
+                # If the argument is a Sample object, splitting is not supported
+                if isinstance(args[0], Samples):
+                    raise ValueError(
+                        "The "
+                        + map_name.lower()
+                        + f" input is specified by a Samples object that cannot be split into multiple arguments corresponding to the non_default_args {non_default_args}."
+                    )
+                split_succeeded, split_args = self._is_stacked_args(*args, is_par=is_par)
+                if split_succeeded:
+                    args = split_args
+                else:
+                    appending_error_message = (
+                        " Additionally, the "
+                        + map_name.lower()
+                        + f" input is specified by a single argument that cannot be split into multiple arguments matching the expected non_default_args {non_default_args}."
+                    )
+
+            # Check if the number of args does not match the number of
+            # non_default_args of the model
+            if len(args) != len(non_default_args):
+                raise ValueError(
+                    "The number of positional arguments does not match the number of non-default arguments of the "
+                    + map_name.lower()
+                    + "."
+                    + appending_error_message
+                )
+
             # Add args to kwargs following the order of non_default_args
             for idx, arg in enumerate(args):
-                kwargs[self._non_default_args[idx]] = arg
+                kwargs[non_default_args[idx]] = arg
+
+        # Check kwargs matches non_default_args
+        if set(list(kwargs.keys())) != set(non_default_args):
+            if map_name == "gradient":
+                error_msg = f"The gradient input is specified by a direction and keywords arguments {list(kwargs.keys())} that does not match the non_default_args of the model {non_default_args}."
+            else:
+                error_msg = (
+                    "The "
+                    + map_name.lower()
+                    + f" input is specified by a keywords arguments {list(kwargs.keys())} that does not match the non_default_args of the "
+                    + map_name
+                    + f" {non_default_args}."
+                )
+
+            raise ValueError(error_msg)
+
+        # Make sure order of kwargs is the same as non_default_args
+        kwargs = {k: kwargs[k] for k in non_default_args}
 
         return kwargs
-        
+
+    def _is_stacked_args(self, *args, is_par=True):
+        """Private function that checks if the input arguments are stacked
+        and splits them if they are."""
+        # Length of args should be 1 if the input is stacked (no partial
+        # stacking is supported)
+        if len(args) > 1:
+            return False, args
+
+        # Type of args should be parameter
+        if not is_par:
+            return False, args
+
+        # args[0] should be numpy array or CUQIarray
+        is_CUQIarray = isinstance(args[0], CUQIarray)
+        is_numpy_array = isinstance(args[0], np.ndarray)
+        if not is_CUQIarray and not is_numpy_array:
+            return False, args
+
+        # Shape of args[0] should be (domain_dim,)
+        if not args[0].shape == (self.domain_dim,):
+            return False, args
+
+        # Ensure domain geometry is _ProductGeometry
+        if not isinstance(
+            self.domain_geometry, cuqi.experimental.geometry._ProductGeometry
+        ):
+            return False, args
+
+        # Split the stacked input
+        split_args = np.split(args[0], self.domain_geometry.stacked_par_split_indices)
+
+        # Covert split args to CUQIarray if input is CUQIarray
+        if is_CUQIarray:
+            split_args = [
+                CUQIarray(arg, is_par=True, geometry=self.domain_geometry.geometries[i])
+                for i, arg in enumerate(split_args)
+            ]
+
+        return True, split_args
+
     def forward(self, *args, is_par=True, **kwargs):
         """ Forward function of the model.
         
-        Forward converts the input to function values (if needed) using the domain geometry of the model.
-        Forward converts the output function values to parameters using the range geometry of the model.
+        Forward converts the input to function values (if needed) using the domain geometry of the model. Then it applies the forward operator to the function values and converts the output to parameters using the range geometry of the model.
 
         Parameters
         ----------
-        *args : ndarray or cuqi.array.CUQIarray
-            The model input.
+        *args : ndarrays or cuqi.array.CUQIarray objects or cuqi.samples.Samples objects
+            Positional arguments for the forward operator. The forward operator input can be specified as either positional arguments or keyword arguments but not both.
 
-        is_par : bool
-            If True the input is assumed to be parameters.
-            If False the input is assumed to be function values.
+            If the input is specified as positional arguments, the order of the arguments should match the non_default_args of the model.
+
+        is_par : bool or a tuple of bools
+            If True, the inputs in `args` or `kwargs` are assumed to be parameters.
+            If False, the inputs in `args` or `kwargs` are assumed to be function values.
+            If `is_par` is a tuple of bools, the inputs are assumed to be parameters or function values based on the corresponding boolean value in the tuple.
         
-        **kwargs : keyword arguments for model input.
-            Keywords must match the names of the non_default_args of the model.
+        **kwargs : keyword arguments
+            keyword arguments for the forward operator. The forward operator input can be specified as either positional arguments or keyword arguments but not both.
+
+            If the input is specified as keyword arguments, the keys should match the non_default_args of the model.
 
         Returns
         -------
-        ndarray or cuqi.array.CUQIarray
+        ndarray or cuqi.array.CUQIarray or cuqi.samples.Samples object
             The model output. Always returned as parameters.
         """
 
-        kwargs = self._parse_args_add_to_kwargs(*args, **kwargs)
-
-        # Check kwargs matches non_default_args
-        if set(list(kwargs.keys())) != set(self._non_default_args):
-            raise ValueError(f"The model input is specified by a keywords arguments {kwargs.keys()} that does not match the non_default_args of the model {self._non_default_args}.")
-
-        # For now only support one input
-        if len(kwargs) > 1:
-            raise ValueError("The model input is specified by more than one argument. This is not supported.")
+        # Add args to kwargs and ensure the order of the arguments matches the
+        # non_default_args of the forward function
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, **kwargs, is_par=is_par, map_name="model"
+        )
 
-        # Get input matching the non_default_args
-        x = kwargs[self._non_default_args[0]]
+        # extract args from kwargs
+        args = list(kwargs.values())
 
-        # If input is a distribution, we simply change the parameter name of model to match the distribution name
-        if isinstance(x, cuqi.distribution.Distribution):
-            if x.dim != self.domain_dim:
-                raise ValueError("Attempting to match parameter name of Model with given distribution, but distribution dimension does not match model domain dimension.")
-            new_model = copy(self)
-            new_model._non_default_args = [x.name] # Defaults to x if distribution had no name
-            return new_model
+        # If input is a distribution, we simply change the parameter name of
+        # model to match the distribution name
+        if all(isinstance(x, cuqi.distribution.Distribution)
+               for x in kwargs.values()):
+            return self._handle_case_when_model_input_is_distributions(kwargs)
 
         # If input is a random variable, we handle it separately
-        if isinstance(x, cuqi.experimental.algebra.RandomVariable):
-            return self._handle_random_variable(x)
-        
+        elif all(isinstance(x, cuqi.experimental.algebra.RandomVariable)
+               for x in kwargs.values()):
+            return self._handle_case_when_model_input_is_random_variables(kwargs)
+
         # If input is a Node from internal abstract syntax tree, we let the Node handle the operation
         # We use NotImplemented to indicate that the operation is not supported from the Model class
         # in case of operations such as "@" that can be interpreted as both __matmul__ and __rmatmul__
         # the operation may be delegated to the Node class.
-        if isinstance(x, cuqi.experimental.algebra.Node):
+        elif any(isinstance(args_i, cuqi.experimental.algebra.Node) for args_i in args):
             return NotImplemented
 
         # Else we apply the forward operator
-        return self._apply_func(self._forward_func,
-                                self.range_geometry,
-                                self.domain_geometry,
-                                x, is_par)
+        # if model has _original_non_default_args, we use it to replace the
+        # kwargs keys so that it matches self._forward_func signature
+        if hasattr(self, '_original_non_default_args'):
+            kwargs = {k:v for k,v in zip(self._original_non_default_args, args)}
+        return self._apply_func(func=self._forward_func,
+                                fwd=True,
+                                is_par=is_par,
+                                **kwargs)
+
+    def _correct_distribution_dimension(self, distributions):
+        """Private function that checks if the dimension of the
+        distributions matches the domain dimension of the model."""
+        if len(distributions) == 1:
+            return list(distributions)[0].dim == self.domain_dim
+        elif len(distributions) > 1 and isinstance(
+            self.domain_geometry, cuqi.experimental.geometry._ProductGeometry
+        ):
+            return all(
+                d.dim == self.domain_geometry.par_dim_list[i]
+                for i, d in enumerate(distributions)
+            )
+        else:
+            return False
 
-    def __call__(self, *args, **kwargs):
-        return self.forward(*args, **kwargs)
+    def _handle_case_when_model_input_is_distributions(self, kwargs):
+        """Private function that handles the case of the input being a
+        distribution or multiple distributions."""
+
+        if not self._correct_distribution_dimension(kwargs.values()):
+            raise ValueError(
+                "Attempting to match parameter name of Model with given distribution(s), but distribution(s) dimension(s) does not match model input dimension(s)."
+            )
+        new_model = copy(self)
+
+        # Store the original non_default_args of the model
+        new_model._original_non_default_args = self._non_default_args
+
+        # Update the non_default_args of the model to match the distribution
+        # names. Defaults to x in the case of only one distribution that has no
+        # name
+        new_model._stored_non_default_args = [x.name for x in kwargs.values()]
+
+        # If there is a repeated name, raise an error
+        if len(set(new_model._stored_non_default_args)) != len(
+            new_model._stored_non_default_args
+        ):
+            raise ValueError(
+                "Attempting to match parameter name of Model with given distributions, but distribution names are not unique. Please provide unique names for the distributions."
+            )
+
+        return new_model
+
+    def _handle_case_when_model_input_is_random_variables(self, kwargs):
+        """ Private function that handles the case of the input being a random variable. """
+        # If random variable is not a leaf-type node (e.g. internal node) we return NotImplemented
+        if any(not isinstance(x.tree, cuqi.experimental.algebra.VariableNode) for x in kwargs.values()):
+            return NotImplemented        
+
+        # Extract the random variable distributions and check dimensions consistency with domain geometry
+        distributions = [value.distribution for value in kwargs.values()]
+        if not self._correct_distribution_dimension(distributions):
+            raise ValueError("Attempting to match parameter name of Model with given random variable(s), but random variable dimension(s) does not match model input dimension(s).")
+
+        new_model = copy(self)
+
+        # Store the original non_default_args of the model
+        new_model._original_non_default_args = self._non_default_args
+
+        # Update the non_default_args of the model to match the random variable
+        # names. Defaults to x in the case of only one random variable that has
+        # no name
+        new_model._stored_non_default_args = [x.name for x in distributions]
+
+        # If there is a repeated name, raise an error
+        if len(set(new_model._stored_non_default_args)) != len(
+            new_model._stored_non_default_args
+        ):
+            raise ValueError(
+                "Attempting to match parameter name of Model with given random variables, but random variables names are not unique. Please provide unique names for the random variables."
+            )
+
+        return new_model
 
-    def gradient(self, direction, wrt, is_direction_par=True, is_wrt_par=True):
-        """ Gradient of the forward operator (Direction-Jacobian product)
+    def gradient(
+        self, direction, *args, is_direction_par=True, is_var_par=True, **kwargs
+    ):
+        """Gradient of the forward operator (Direction-Jacobian product)
 
-        For non-linear models the gradient is computed using the
-        forward operator and the Jacobian of the forward operator.
+        The gradient computes the Vector-Jacobian product (VJP) of the forward operator evaluated at the given model input and the given vector (direction).
 
         Parameters
         ----------
-        direction : ndarray
-            The direction to compute the gradient. The Jacobian is applied to this direction.
+        direction : ndarray or cuqi.array.CUQIarray
+            The direction at which to compute the gradient.
 
-        wrt : ndarray
-            The point to compute the Jacobian at. This is only used for non-linear models.
+        *args : ndarrays or cuqi.array.CUQIarray objects
+            Positional arguments for the values at which to compute the gradient. The gradient operator input can be specified as either positional arguments or keyword arguments but not both.
+
+            If the input is specified as positional arguments, the order of the arguments should match the non_default_args of the model.
 
         is_direction_par : bool
             If True, `direction` is assumed to be parameters.
             If False, `direction` is assumed to be function values.
 
-        is_wrt_par : bool
-            If True, `wrt` is assumed to be parameters.
-            If False, `wrt` is assumed to be function values.
-        
+        is_var_par : bool or a tuple of bools
+            If True, the inputs in `args` or `kwargs` are assumed to be parameters.
+            If False, the inputs in `args` or `kwargs` are assumed to be function values.
+            If `is_var_par` is a tuple of bools, the inputs in `args` or `kwargs` are assumed to be parameters or function values based on the corresponding boolean value in the tuple.
         """
-        # Obtain the parameters representation of wrt and raise an error if it
-        # cannot be obtained
-        error_message = \
-            "For the gradient to be computed, is_wrt_par needs " +\
-            "to be True and wrt needs to be parameter value, not function " +\
-            "value. Alternatively, the model domain_geometry: "+\
-            f"{self.domain_geometry} " +\
-            "should have an implementation of the method fun2par"
+        # Add args to kwargs and ensure the order of the arguments matches the
+        # non_default_args of the forward function
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, **kwargs, is_par=is_var_par, map_name="gradient"
+        )
+
+        # Obtain the parameters representation of the variables and raise an
+        # error if it cannot be obtained
+        error_message = (
+            "For the gradient to be computed, is_var_par needs to be True and the variables in kwargs needs to be parameter value, not function value. Alternatively, the model domain_geometry:"
+            + f" {self.domain_geometry} "
+            + "should have an implementation of the method fun2par"
+        )
         try:
-            wrt_par = self._2par(wrt, 
-                                 geometry=self.domain_geometry,
-                                 is_par=is_wrt_par,
-                                 to_CUQIarray=False,
-                                 )
+            kwargs_par = self._2par(
+                geometry=self.domain_geometry,
+                is_par=is_var_par,
+                to_CUQIarray=False,
+                **kwargs,
+            )
         # NotImplementedError will be raised if fun2par of the geometry is not
         # implemented and ValueError will be raised when imap is not set in
         # MappedGeometry
         except ValueError as e:
-            raise ValueError(error_message +
-                             " ,including an implementation of imap for " +
-                             "MappedGeometry")
+            raise ValueError(
+                error_message
+                + " ,including an implementation of imap for MappedGeometry"
+            )
         except NotImplementedError as e:
             raise NotImplementedError(error_message)
-        
-        # Check for other errors that may prevent computing the gradient
-        self._check_gradient_can_be_computed(direction, wrt)
-
-        wrt = self._2fun(wrt, self.domain_geometry, is_par=is_wrt_par)
-
-        # Store if the input direction is CUQIarray
-        is_direction_CUQIarray = type(direction) is CUQIarray
 
-        direction = self._2fun(direction,
-                               self.range_geometry,
-                               is_par=is_direction_par)
-
-        grad = self._gradient_func(direction, wrt)
-        grad_is_par = False # Assume gradient is function values
-        
-        # If domain_geometry has gradient attribute, we apply it to the gradient
-        # The gradient returned by the domain_geometry.gradient is assumed to be
-        # parameters
-        if hasattr(self.domain_geometry, 'gradient'):
-            grad = self.domain_geometry.gradient(grad, wrt_par)
-            grad_is_par = True # Gradient is parameters
-
-        # we convert the computed gradient to parameters
-        grad = self._2par(grad,
-                          self.domain_geometry,
-                          to_CUQIarray=is_direction_CUQIarray,
-                          is_par=grad_is_par)
+        # Check for other errors that may prevent computing the gradient
+        self._check_gradient_can_be_computed(direction, kwargs)
+
+        # Also obtain the function values representation of the variables
+        kwargs_fun = self._2fun(
+            geometry=self.domain_geometry, is_par=is_var_par, **kwargs
+        )
+
+        # Store if any of the inputs is a CUQIarray
+        to_CUQIarray = isinstance(direction, CUQIarray) or any(
+            isinstance(x, CUQIarray) for x in kwargs_fun.values()
+        )
+
+        # Turn to_CUQIarray to a tuple of bools of the same length as kwargs_fun
+        to_CUQIarray = tuple([to_CUQIarray] * len(kwargs_fun))
+
+        # Convert direction to function value
+        direction_fun = self._2fun(
+            direction=direction, geometry=self.range_geometry, is_par=is_direction_par
+        )
+
+        # If model has _original_non_default_args, we use it to replace the
+        # kwargs keys so that it matches self._gradient_func signature
+        if hasattr(self, '_original_non_default_args'):
+            args_fun = list(kwargs_fun.values())
+            kwargs_fun = {
+                k: v for k, v in zip(self._original_non_default_args, args_fun)
+            }
+        # Append the direction to the kwargs_fun as first input
+        kwargs_fun_grad_input = {**direction_fun, **kwargs_fun}
+
+        # Form 1 of gradient (callable)
+        if callable(self._gradient_func):
+            grad = self._gradient_func(**kwargs_fun_grad_input)
+            grad_is_par = False  # Assume gradient is function value
+
+        # Form 2 of gradient (tuple of callables)
+        elif isinstance(self._gradient_func, tuple):
+            grad = []
+            for i, grad_func in enumerate(self._gradient_func):
+                if grad_func is not None:
+                    grad.append(grad_func(**kwargs_fun_grad_input))
+                else:
+                    grad.append(None)
+                    # set the ith item of to_CUQIarray tuple to False
+                    # because the ith gradient is None
+                    to_CUQIarray = to_CUQIarray[:i] + (False,) + to_CUQIarray[i + 1 :]
+            grad_is_par = False  # Assume gradient is function value
+
+        grad = self._apply_chain_rule_to_account_for_domain_geometry_gradient(
+            kwargs_par, grad, grad_is_par, to_CUQIarray
+        )
+
+        if len(grad) == 1:
+            return list(grad.values())[0]
+        elif self._gradient_output_stacked:
+            return np.hstack(
+                [
+                    (
+                        v.to_numpy()
+                        if isinstance(v, CUQIarray)
+                        else force_ndarray(v, flatten=True)
+                    )
+                    for v in list(grad.values())
+                ]
+            )
 
         return grad
-    
-    def _check_gradient_can_be_computed(self, direction, wrt):
-        """ Private function that checks if the gradient can be computed. By
+
+    def _check_gradient_can_be_computed(self, direction, kwargs_dict):
+        """Private function that checks if the gradient can be computed. By
         raising an error for the cases where the gradient cannot be computed."""
 
         # Raise an error if _gradient_func function is not set
         if self._gradient_func is None:
             raise NotImplementedError("Gradient is not implemented for this model.")
-        
-        # Raise error if either the direction or wrt are Samples object
-        if isinstance(direction, Samples) or isinstance(wrt, Samples):
-            raise ValueError("cuqi.samples.Samples input values for arguments `direction` and `wrt` are not supported")
-        
+
+        # Raise an error if either the direction or kwargs are Samples objects
+        if isinstance(direction, Samples) or any(
+            isinstance(x, Samples) for x in kwargs_dict.values()
+        ):
+            raise NotImplementedError(
+                "Gradient is not implemented for input of type Samples."
+            )
+
         # Raise an error if range_geometry is not in the list returned by
-        # `_get_identity_geometries()`. i.e. The Jacobian of its 
-        # par2fun map is not identity.  
-        #TODO: Add range geometry gradient to the chain rule 
+        # `_get_identity_geometries()`. i.e. The Jacobian of its
+        # par2fun map is not identity.
+        # TODO: Add range geometry gradient to the chain rule
         if not type(self.range_geometry) in _get_identity_geometries():
-            raise NotImplementedError("Gradient not implemented for model {} with range geometry {}".format(self,self.range_geometry)) 
-        
-        # Raise an error if domain_geometry does not have gradient attribute and
-        # is not in the list returned by `_get_identity_geometries()`. i.e. the
-        # Jacobian of its par2fun map is not identity.  
-        if not hasattr(self.domain_geometry, 'gradient') and \
-            not type(self.domain_geometry) in _get_identity_geometries():
-            raise NotImplementedError("Gradient not implemented for model {} with domain geometry {}".format(self,self.domain_geometry))
-
-    def _handle_random_variable(self, x):
-        """ Private function that handles the case of the input being a random variable. """
-        # If random variable is not a leaf-type node (e.g. internal node) we return NotImplemented
-        if not isinstance(x.tree, cuqi.experimental.algebra.VariableNode):
-            return NotImplemented        
-        
-        # In leaf-type node case we simply change the parameter name of model to match the random variable name
-        dist = x.distribution
-        if dist.dim != self.domain_dim:
-            raise ValueError("Attempting to match parameter name of Model with given random variable, but random variable dimension does not match model domain dimension.")
-        
-        new_model = copy(self)
-        new_model._non_default_args = [dist.name]
-        return new_model
+            raise NotImplementedError(
+                "Gradient is not implemented for model {} with range geometry {}. You can use one of the geometries in the list {}.".format(
+                    self,
+                    self.range_geometry,
+                    [i_g.__name__ for i_g in _get_identity_geometries()],
+                )
+            )
+
+        # Raise an error if domain_geometry (or its components in case of
+        # _ProductGeometry) does not have gradient attribute and is not in the
+        # list returned by `_get_identity_geometries()`. i.e. The Jacobian of its
+        # par2fun map is not identity.
+        domain_geometries = (
+            self.domain_geometry.geometries
+            if isinstance(
+                self.domain_geometry, cuqi.experimental.geometry._ProductGeometry
+            )
+            else [self.domain_geometry]
+        )
+        for domain_geometry in domain_geometries:
+            if (
+                not hasattr(domain_geometry, "gradient")
+                and not type(domain_geometry) in _get_identity_geometries()
+            ):
+                raise NotImplementedError(
+                    "Gradient is not implemented for model \n{}\nwith domain geometry (or domain geometry component) {}. The domain geometries should have gradient method or be from the geometries in the list {}.".format(
+                        self,
+                        domain_geometry,
+                        [i_g.__name__ for i_g in _get_identity_geometries()],
+                    )
+                )
+
+    def _apply_chain_rule_to_account_for_domain_geometry_gradient(self,
+                                                                  kwargs_par,
+                                                                  grad,
+                                                                  grad_is_par,
+                                                                  to_CUQIarray):
+        """ Private function that applies the chain rule to account for the
+        gradient of the domain geometry. That is, it computes the gradient of
+        the function values with respect to the parameters values."""
+        # Create list of domain geometries
+        geometries = (
+            self.domain_geometry.geometries
+            if isinstance(self.domain_geometry, cuqi.experimental.geometry._ProductGeometry)
+            else [self.domain_geometry]
+        )
+
+        # turn grad_is_par to a tuple of bools if it is not already
+        if isinstance(grad_is_par, bool):
+            grad_is_par = tuple([grad_is_par]*len(grad))
+
+        # If the domain geometry is a _ProductGeometry and the gradient is
+        # stacked, split it
+        if (
+            isinstance(
+                self.domain_geometry, cuqi.experimental.geometry._ProductGeometry
+            )
+            and not isinstance(grad, (list, tuple))
+            and isinstance(grad, np.ndarray)
+        ):
+            grad = np.split(grad, self.domain_geometry.stacked_par_split_indices)
+
+        # If the domain geometry is not a _ProductGeometry, turn grad into a
+        # list of length 1, so that we can iterate over it
+        if not isinstance(self.domain_geometry, cuqi.experimental.geometry._ProductGeometry):
+            grad = [grad]
+
+        # apply the gradient of each geometry component
+        grad_kwargs = {}
+        for i, (k, v_par) in enumerate(kwargs_par.items()):
+            if hasattr(geometries[i], 'gradient') and grad[i] is not None:
+                grad_kwargs[k] = geometries[i].gradient(grad[i], v_par)
+                # update the ith component of grad_is_par to True
+                grad_is_par = grad_is_par[:i] + (True,) + grad_is_par[i+1:]
+            else:
+                grad_kwargs[k] = grad[i]
+
+        # convert the computed gradient to parameters
+        grad = self._2par(geometry=self.domain_geometry,
+                          to_CUQIarray=to_CUQIarray,
+                          is_par=grad_is_par,
+                          **grad_kwargs)
+
+        return grad
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
 
     def __len__(self):
         return self.range_dim
 
     def __repr__(self) -> str:
-        return "CUQI {}: {} -> {}.\n    Forward parameters: {}.".format(self.__class__.__name__,self.domain_geometry,self.range_geometry,cuqi.utilities.get_non_default_args(self))
-
+        kwargs = {}
+        if self.number_of_inputs > 1:
+            pad = " " * len("CUQI {}: ".format(self.__class__.__name__))
+            kwargs["pad"]=pad  
+        return "CUQI {}: {} -> {}.\n    Forward parameters: {}.".format(self.__class__.__name__,self.domain_geometry.__repr__(**kwargs),self.range_geometry,self._non_default_args)
 
 class AffineModel(Model):
     """ Model class representing an affine model, i.e. a linear operator with a fixed shift. For linear models, represented by a linear operator only, see :class:`~cuqi.model.LinearModel`.
@@ -533,7 +1138,7 @@ def __init__(self, linear_operator, shift, linear_operator_adjoint=None, range_g
         if hasattr(linear_operator, '__matmul__') and hasattr(linear_operator, 'T'):
             if linear_operator_adjoint is not None:
                 raise ValueError("Adjoint of linear operator should not be provided when linear operator is a matrix. If you want to provide an adjoint, use a callable function for the linear operator.")
-            
+
             matrix = linear_operator
 
             linear_operator = lambda x: matrix@x
@@ -559,11 +1164,50 @@ def __init__(self, linear_operator, shift, linear_operator_adjoint=None, range_g
         if linear_operator_adjoint is not None and not callable(linear_operator_adjoint):
             raise TypeError("Linear operator adjoint must be defined as a callable function of some kind")
 
+        # If linear operator is of type Model, it needs to be a LinearModel
+        if isinstance(linear_operator, Model) and not isinstance(
+            linear_operator, LinearModel
+        ):
+            raise TypeError(
+                "The linear operator should be a LinearModel object, a callable function or a matrix."
+            )
+
+        # If the adjoint operator is of type Model, it needs to be a LinearModel
+        if isinstance(linear_operator_adjoint, Model) and not isinstance(
+            linear_operator_adjoint, LinearModel
+        ):
+            raise TypeError(
+                "The adjoint linear operator should be a LinearModel object, a callable function or a matrix."
+            )
+
+        # Additional checks if the linear_operator is not a LinearModel:
+        if not isinstance(linear_operator, LinearModel):
+            # Ensure the linear operator has exactly one input argument
+            if len(cuqi.utilities.get_non_default_args(linear_operator)) != 1:
+                raise ValueError(
+                    "The linear operator should have exactly one input argument."
+                )
+            # Ensure the adjoint linear operator has exactly one input argument
+            if (
+                linear_operator_adjoint is not None
+                and len(cuqi.utilities.get_non_default_args(linear_operator_adjoint))
+                != 1
+            ):
+                raise ValueError(
+                    "The adjoint linear operator should have exactly one input argument."
+                )
+
         # Check size of shift and match against range_geometry
         if not np.isscalar(shift):
             if len(shift) != range_geometry.par_dim:
                 raise ValueError("The shift should have the same dimension as the range geometry.")
 
+        # Store linear operator privately
+        # Note: we need to set the _linear_operator before calling the
+        # super().__init__() because it is needed when calling the property
+        # _non_default_args within the super().__init__()
+        self._linear_operator = linear_operator
+
         # Initialize Model class
         super().__init__(linear_operator, range_geometry, domain_geometry)
 
@@ -573,20 +1217,27 @@ def __init__(self, linear_operator, shift, linear_operator_adjoint=None, range_g
         # Store shift as private attribute
         self._shift = shift
 
-        # Store linear operator privately
-        self._linear_operator = linear_operator
 
         # Store adjoint function
         self._linear_operator_adjoint = linear_operator_adjoint
 
         # Define gradient
-        self._gradient_func = lambda direction, wrt: linear_operator_adjoint(direction)
+        self._gradient_func = lambda direction, *args, **kwargs: linear_operator_adjoint(direction)
 
         # Update forward function to include shift (overwriting the one from Model class)
         self._forward_func = lambda *args, **kwargs: linear_operator(*args, **kwargs) + shift
 
-        # Use arguments from user's callable linear operator (overwriting those found by Model class)
-        self._non_default_args = cuqi.utilities.get_non_default_args(linear_operator)
+        # Set stored_non_default_args to None
+        self._stored_non_default_args = None
+
+    @property
+    def _non_default_args(self):
+        if self._stored_non_default_args is None:
+            # Use arguments from user's callable linear operator
+            self._stored_non_default_args = cuqi.utilities.get_non_default_args(
+                self._linear_operator
+            )
+        return self._stored_non_default_args
 
     @property
     def shift(self):
@@ -599,19 +1250,35 @@ def shift(self, value):
         self._shift = value
         self._forward_func = lambda *args, **kwargs: self._linear_operator(*args, **kwargs) + value
 
-    def _forward_func_no_shift(self, x, is_par=True):
-        """ Helper function for computing the forward operator without the shift. """
-        return self._apply_func(self._linear_operator,
-                self.range_geometry,
-                self.domain_geometry,
-                x, is_par)
+    def _forward_func_no_shift(self, *args, is_par=True, **kwargs):
+        """Helper function for computing the forward operator without the shift."""
+        # convert args to kwargs
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, **kwargs, map_name="model", is_par=is_par
+        )
+        args = list(kwargs.values())
+        # if model has _original_non_default_args, we use it to replace the
+        # kwargs keys so that it matches self._linear_operator signature
+        if hasattr(self, '_original_non_default_args'):
+            kwargs = {k:v for k,v in zip(self._original_non_default_args, args)}
+        return self._apply_func(self._linear_operator, **kwargs, is_par=is_par)
+
+    def _adjoint_func_no_shift(self, *args, is_par=True, **kwargs):
+        """Helper function for computing the adjoint operator without the shift."""
+        # convert args to kwargs
+        kwargs = self._parse_args_add_to_kwargs(
+            *args,
+            **kwargs,
+            map_name='adjoint',
+            is_par=is_par,
+            non_default_args=cuqi.utilities.get_non_default_args(
+                self._linear_operator_adjoint
+            ),
+        )
+        return self._apply_func(
+            self._linear_operator_adjoint, **kwargs, is_par=is_par, fwd=False
+        )
 
-    def _adjoint_func_no_shift(self, y, is_par=True):
-        """ Helper function for computing the adjoint operator without the shift. """
-        return self._apply_func(self._linear_operator_adjoint,
-                self.domain_geometry,
-                self.range_geometry,
-                y, is_par)
 
 class LinearModel(AffineModel):
     """Model based on a Linear forward operator.
@@ -677,50 +1344,67 @@ def adjoint(y):
     Note that you would need to specify the range and domain geometries in this
     case as they cannot be inferred from the forward and adjoint functions.
     """
-    
+
     def __init__(self, forward, adjoint=None, range_geometry=None, domain_geometry=None):
 
-        #Initialize as AffineModel with shift=0
+        # Initialize as AffineModel with shift=0
         super().__init__(forward, 0, adjoint, range_geometry, domain_geometry)
 
-    def adjoint(self, y, is_par=True):
+    def adjoint(self, *args, is_par=True, **kwargs):
         """ Adjoint of the model.
         
-        Adjoint converts the input to function values (if needed) using the range geometry of the model.
-        Adjoint converts the output function values to parameters using the range geometry of the model.
+        Adjoint converts the input to function values (if needed) using the range geometry of the model then applies the adjoint operator to the function values and converts the output function values to parameters using the domain geometry of the model.
 
         Parameters
         ----------
-        y : ndarray or cuqi.array.CUQIarray
-            The adjoint model input.
+        *args : ndarrays or cuqi.array.CUQIarray object
+            Positional arguments for the adjoint operator ( maximum one argument). The adjoint operator input can be specified as either positional arguments or keyword arguments but not both.
+
+        **kwargs : keyword arguments
+            keyword arguments for the adjoint operator (maximum one argument). The adjoint operator input can be specified as either positional arguments or keyword arguments but not both.
+
+            If the input is specified as keyword arguments, the keys should match the non_default_args of the model.
 
         Returns
         -------
         ndarray or cuqi.array.CUQIarray
             The adjoint model output. Always returned as parameters.
         """
+        kwargs = self._parse_args_add_to_kwargs(
+            *args,
+            **kwargs,
+            map_name='adjoint',
+            is_par=is_par,
+            non_default_args=cuqi.utilities.get_non_default_args(
+                self._linear_operator_adjoint
+            ),
+        )
+
+        # length of kwargs should be 1
+        if len(kwargs) > 1:
+            raise ValueError(
+                "The adjoint operator input is specified by more than one argument. This is not supported."
+            )
         if self._linear_operator_adjoint is None:
             raise ValueError("No adjoint operator was provided for this model.")
-        return self._apply_func(self._linear_operator_adjoint,
-                                self.domain_geometry,
-                                self.range_geometry,
-                                y, is_par)
+        return self._apply_func(
+            self._linear_operator_adjoint, **kwargs, is_par=is_par, fwd=False
+        )
+
+    def __matmul__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
 
-    def __matmul__(self, x):
-        return self.forward(x)
-        
     def get_matrix(self):
         """
         Returns an ndarray with the matrix representing the forward operator.
         """
-
         if self._matrix is not None: #Matrix exists so return it
             return self._matrix
         else:
-            #TODO: Can we compute this faster while still in sparse format?
+            # TODO: Can we compute this faster while still in sparse format?
             mat = csc_matrix((self.range_dim,0)) #Sparse (m x 1 matrix)
             e = np.zeros(self.domain_dim)
-            
+
             # Stacks sparse matrices on csc matrix
             for i in range(self.domain_dim):
                 e[i] = 1
@@ -728,7 +1412,7 @@ def get_matrix(self):
                 mat = hstack((mat,col_vec[:,None])) #mat[:,i] = self.forward(e)
                 e[i] = 0
 
-            #Store matrix for future use
+            # Store matrix for future use
             self._matrix = mat
 
             return self._matrix   
@@ -736,26 +1420,31 @@ def get_matrix(self):
     @property
     def T(self):
         """Transpose of linear model. Returns a new linear model acting as the transpose."""
-        transpose = LinearModel(self.adjoint, self.forward, self.domain_geometry, self.range_geometry)
+        transpose = LinearModel(
+            self._linear_operator_adjoint,
+            self._linear_operator,
+            self.domain_geometry,
+            self.range_geometry,
+        )
         if self._matrix is not None:
             transpose._matrix = self._matrix.T
         return transpose
-        
+
 
 class PDEModel(Model):
     """
     Model based on an underlying cuqi.pde.PDE.
-    In the forward operation the PDE is assembled, solved and observed.
+    In the forward method the PDE is assembled, solved and observed.
     
     Parameters
     -----------
-    forward : 2D ndarray or callable function.
-        Forward operator assembling, solving and observing the pde.
+    PDE : cuqi.pde.PDE
+        The PDE that specifies the forward operator.
 
-    range_geometry : integer or cuqi.geometry.Geometry (optional)
+    range_geometry : integer or cuqi.geometry.Geometry, optional
         If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
 
-    domain_geometry : integer or cuqi.geometry.Geometry (optional)
+    domain_geometry : integer or cuqi.geometry.Geometry, optional
         If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
 
 
@@ -766,21 +1455,37 @@ def __init__(self, PDE: cuqi.pde.PDE, range_geometry, domain_geometry):
 
         if not isinstance(PDE, cuqi.pde.PDE):
             raise ValueError("PDE needs to be a cuqi PDE.")
+        # PDE needs to be set before calling super().__init__
+        # for the property _non_default_args to work
+        self.pde = PDE
+        self._stored_non_default_args = None
 
-        super().__init__(self._forward_func, range_geometry, domain_geometry, gradient=self._gradient_func)
+        super().__init__(self._forward_func, range_geometry, domain_geometry)
 
-        self.pde = PDE
+    @property
+    def _non_default_args(self):
+        if self._stored_non_default_args is None:
+            # extract the non-default arguments of the PDE
+            self._stored_non_default_args = cuqi.utilities.get_non_default_args(
+                self.pde.PDE_form
+            )
+            # remove t from the non-default arguments
+            self._stored_non_default_args = self._non_default_args
+            if "t" in self._non_default_args:
+                self._stored_non_default_args.remove("t")
 
-    def _forward_func(self, x):
-        
-        self.pde.assemble(parameter=x)
+        return self._stored_non_default_args
+
+    def _forward_func(self, **kwargs):
+
+        self.pde.assemble(**kwargs)
 
         sol, info = self.pde.solve()
 
         obs = self.pde.observe(sol)
 
         return obs
-    
+
     def _gradient_func(self, direction, wrt):
         """ Compute direction-Jacobian product (gradient) of the model. """
         if hasattr(self.pde, "gradient_wrt_parameter"):
@@ -793,4 +1498,3 @@ def _gradient_func(self, direction, wrt):
     # Add the underlying PDE class name to the repr.
     def __repr__(self) -> str:
         return super().__repr__()+"\n    PDE: {}.".format(self.pde.__class__.__name__)
-        
diff --git a/cuqi/pde/_pde.py b/cuqi/pde/_pde.py
index 1c47b8a74..a5ff95f8a 100644
--- a/cuqi/pde/_pde.py
+++ b/cuqi/pde/_pde.py
@@ -4,6 +4,7 @@
 from scipy.interpolate import interp1d
 import numpy as np
 
+
 class PDE(ABC):
     """
     Parametrized PDE abstract base class
@@ -30,7 +31,7 @@ def __init__(self, PDE_form, grid_sol=None, grid_obs=None, observation_map=None)
         self.observation_map = observation_map
 
     @abstractmethod
-    def assemble(self,parameter):
+    def assemble(self, *args, **kwargs):
         pass
 
     @abstractmethod
@@ -155,9 +156,9 @@ class SteadyStateLinearPDE(LinearPDE):
     def __init__(self, PDE_form, **kwargs):
         super().__init__(PDE_form, **kwargs)
 
-    def assemble(self, parameter):
+    def assemble(self, *args, **kwargs):
         """Assembles differential operator and rhs according to PDE_form"""
-        self.diff_op, self.rhs = self.PDE_form(parameter)
+        self.diff_op, self.rhs = self.PDE_form(*args, **kwargs)
 
     def solve(self):
         """Solve the PDE and returns the solution and an information variable `info` which is a tuple of all variables returned by the function `linalg_solve` after the solution."""
@@ -178,7 +179,7 @@ def observe(self, solution):
             solution_obs = self.observation_map(solution_obs)
                 
         return solution_obs
-        
+
 class TimeDependentLinearPDE(LinearPDE):
     """Time Dependent Linear PDE with fixed time stepping using Euler method (backward or forward).
     
@@ -234,13 +235,16 @@ def method(self, value):
                 "method can be set to either `forward_euler` or `backward_euler`")
         self._method = value
 
-    def assemble(self, parameter):
+    def assemble(self, *args, **kwargs):
         """Assemble PDE"""
-        self._parameter = parameter
+        self._parameter_kwargs = kwargs
+        self._parameter_args = args
 
     def assemble_step(self, t):
         """Assemble time step at time t"""
-        self.diff_op, self.rhs, self.initial_condition = self.PDE_form(self._parameter, t)
+        self.diff_op, self.rhs, self.initial_condition = self.PDE_form(
+            *self._parameter_args, **self._parameter_kwargs, t=t
+        )
 
     def solve(self):
         """Solve PDE by time-stepping"""
@@ -279,7 +283,7 @@ def observe(self, solution):
         # Interpolate solution in time and space to the observation
         # time and space
         else:
-            # Raise error if solution is 2D or 3D in space 
+            # Raise error if solution is 2D or 3D in space
             if len(solution.shape) > 2:
                 raise ValueError("Interpolation of solutions of 2D and 3D "+ 
                                  "space dimensions based on the provided "+
@@ -287,7 +291,7 @@ def observe(self, solution):
                                  "You can, instead, pass a custom "+
                                  "observation_map and pass grid_obs and "+
                                  "time_obs as None.")
-            
+
             # Interpolate solution in space and time to the observation
             # time and space
             solution_obs = scipy.interpolate.RectBivariateSpline(
@@ -297,7 +301,7 @@ def observe(self, solution):
         # Apply observation map
         if self.observation_map is not None:
             solution_obs = self.observation_map(solution_obs)
-        
+
         # squeeze if only one time observation
         if len(self._time_obs) == 1:
             solution_obs = solution_obs.squeeze()
diff --git a/cuqi/problem/_problem.py b/cuqi/problem/_problem.py
index 3323214c6..9d52bc4d2 100644
--- a/cuqi/problem/_problem.py
+++ b/cuqi/problem/_problem.py
@@ -304,6 +304,7 @@ def sample_posterior(self, Ns, Nb=None, callback=None, experimental=False) -> cu
             The signature of the callback function is `callback(sample, sample_index)`,
             where `sample` is the current sample and `sample_index` is the index of the sample.
             An example is shown in demos/demo31_callback.py.
+            Note: if the parameter `experimental` is set to True, the callback function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature in the case is: `callback(sampler, sample_index, num_of_samples)`.
 
         experimental : bool, *Optional*
             If set to True, the sampler selection will use the samplers from the :mod:`cuqi.experimental.mcmc` module.
@@ -848,16 +849,14 @@ def _sampleGibbs(self, Ns, Nb, callback=None, experimental=False):
             print(f"burn-in: {Nb/Ns*100:g}%")
             print("")
 
-            if callback is not None:
-                raise NotImplementedError("Callback not implemented for Gibbs sampler")
-
             # Start timing
             ti = time.time()
 
             # Sampling strategy
             sampling_strategy = self._determine_sampling_strategy(experimental=True)
 
-            sampler = cuqi.experimental.mcmc.HybridGibbs(self._target, sampling_strategy)
+            sampler = cuqi.experimental.mcmc.HybridGibbs(
+                self._target, sampling_strategy, callback=callback)
             sampler.warmup(Nb)
             sampler.sample(Ns)
             samples = sampler.get_samples()
@@ -876,7 +875,7 @@ def _sampleGibbs(self, Ns, Nb, callback=None, experimental=False):
             print("")
 
             if callback is not None:
-                raise NotImplementedError("Callback not implemented for Gibbs sampler")
+                raise NotImplementedError("Callback not implemented for Gibbs sampler. It is only implemented for experimental Gibbs sampler.")
 
             # Start timing
             ti = time.time()
diff --git a/cuqi/solver/_solver.py b/cuqi/solver/_solver.py
index 9ceacf025..d1c632273 100644
--- a/cuqi/solver/_solver.py
+++ b/cuqi/solver/_solver.py
@@ -196,8 +196,11 @@ class ScipyLSQ(object):
         'trf', Trust Region Reflective algorithm: for large sparse problems with bounds.
         'dogbox', dogleg algorithm with rectangular trust regions, for small problems with bounds.
         'lm', Levenberg-Marquardt algorithm as implemented in MINPACK. Doesn't handle bounds and sparse Jacobians.
+    tol : The numerical tolerance for convergence checks.
+    maxit : The maximum number of iterations.
+    kwargs : Additional keyword arguments passed to scipy's least_squares. Empty by default. See documentation for scipy.optimize.least_squares
     """
-    def __init__(self, func, x0, jacfun='2-point', method='trf', loss='linear', tol=1e-6, maxit=1e4):
+    def __init__(self, func, x0, jacfun='2-point', method='trf', loss='linear', tol=1e-6, maxit=1e4, **kwargs):
         self.func = func
         self.x0 = x0
         self.jacfun = jacfun
@@ -205,6 +208,7 @@ def __init__(self, func, x0, jacfun='2-point', method='trf', loss='linear', tol=
         self.loss = loss
         self.tol = tol
         self.maxit = int(maxit)
+        self.kwargs = kwargs
     
     def solve(self):
         """Runs optimization algorithm and returns solution and info.
@@ -215,7 +219,7 @@ def solve(self):
             Solution found (array_like) and optimization information (dictionary).
         """
         solution = least_squares(self.func, self.x0, jac=self.jacfun, \
-                                method=self.method, loss=self.loss, xtol=self.tol, max_nfev=self.maxit)
+                                method=self.method, loss=self.loss, xtol=self.tol, max_nfev=self.maxit, **self.kwargs)
         info = {"success": solution['success'],
                 "message": solution['message'],
                 "func": solution['fun'],
diff --git a/cuqi/testproblem/_testproblem.py b/cuqi/testproblem/_testproblem.py
index 4c82c82b9..01ee80e99 100644
--- a/cuqi/testproblem/_testproblem.py
+++ b/cuqi/testproblem/_testproblem.py
@@ -863,10 +863,9 @@ def PDE_form(IC, t): return (Dxx, np.zeros(N), IC)
         # Bayesian model
         x = cuqi.distribution.Gaussian(np.zeros(model.domain_dim), 1)
         y = cuqi.distribution.Gaussian(model(x), sigma2)
-        
-        # Initialize Deconvolution as BayesianProblem problem
-        super().__init__(y, x, y=data)
 
+        # Initialize Heat1D as BayesianProblem problem
+        super().__init__(y, x, y=data)
         # Store exact values
         self.exactSolution = x_exact
         self.exactData = y_exact
diff --git a/cuqi/utilities/__init__.py b/cuqi/utilities/__init__.py
index 602038f3a..5a7874170 100644
--- a/cuqi/utilities/__init__.py
+++ b/cuqi/utilities/__init__.py
@@ -14,8 +14,10 @@
     plot_1D_density,
     plot_2D_density,
     count_nonzero,
+    count_within_bounds,
     count_constant_components_1D,
-    count_constant_components_2D
+    count_constant_components_2D,
+    piecewise_linear_1D_DoF
 )
 
 from ._get_python_variable_name import _get_python_variable_name
diff --git a/cuqi/utilities/_utilities.py b/cuqi/utilities/_utilities.py
index 3e7903632..a246d2865 100644
--- a/cuqi/utilities/_utilities.py
+++ b/cuqi/utilities/_utilities.py
@@ -188,7 +188,7 @@ def approx_derivative(func, wrt, direction=None, epsilon=np.sqrt(np.finfo(float)
     # We compute the Jacobian matrix of func using forward differences.
     # If the function is scalar-valued, we compute the gradient instead.
     # If the direction is provided, we compute the direction-Jacobian product.
-    wrt = np.asfarray(wrt)
+    wrt = np.asarray(wrt)
     f0 = func(wrt)
     Matr = np.zeros([infer_len(wrt), infer_len(f0)])
     dx = np.zeros(len(wrt))
@@ -360,8 +360,29 @@ def count_nonzero(x, threshold = 1e-6):
             Theshold for considering a value as nonzero.
         """
         return np.sum([np.abs(v) >= threshold for v in x])
+
+        
+def count_within_bounds(x, lower_bounds, upper_bounds, threshold = 1e-6, exception = np.nan):
+        """ Returns the number of values in an array whose value lies between the provided lower and upper bounds.
+
+        Parameters
+        ----------
+        x : `np.ndarray` 
+            Array to count elements of.
+
+        lower_bounds : `np.ndarray` 
+            Lower bound on values to disregard when counting.
+
+        upper_bounds : `np.ndarray` 
+            Upper bound on values to disregard.
+
+        threshold : float
+            Theshold for considering a value as nonzero.
+        """
+        return np.sum([l + threshold <= v and v <= u - threshold and not (exception - threshold <= v and v <= exception + threshold) for v, l, u in zip(x, lower_bounds, upper_bounds)])
+        
     
-def count_constant_components_1D(x, threshold = 1e-2, lower = -np.inf, upper = np.inf):
+def count_constant_components_1D(x, threshold = 1e-2, lower = -np.inf, upper = np.inf, exception = np.nan):
         """ Returns the number of piecewise constant components in a one-dimensional array
 
         Parameters
@@ -372,29 +393,40 @@ def count_constant_components_1D(x, threshold = 1e-2, lower = -np.inf, upper = n
         threshold : float
             Strict theshold on when the difference of neighbouring values is considered zero.
 
-        lower : float
+        lower : float or numpy.ndarray
             Piecewise constant components below this value are not counted.
 
-        upper : float
+        upper : float or numpy.ndarray
             Piecewise constant components above this value are not counted.
         """
 
+        if not isinstance(lower, np.ndarray):
+            lower = lower*np.ones_like(x)
+
+        if not isinstance(upper, np.ndarray):
+            upper = upper*np.ones_like(x)
+
         counter = 0
-        if x[0] > lower and x[0] < upper:
+        index = 0
+        if (x[index] > lower[index] and
+            x[index] < upper[index] and
+            not (exception - threshold <= x[index] and x[index] <= exception + threshold)):
             counter += 1
         
         x_previous = x[0]
 
         for x_current in x[1:]:
+            index += 1
             if (abs(x_previous - x_current) >= threshold and
-                x_current > lower and
-                x_current < upper):
+                x_current > lower[index] and
+                x_current < upper[index] and
+                not (exception - threshold <= x_current and x_current <= exception + threshold)):
                     counter += 1
 
             x_previous = x_current
     
         return counter
-        
+
 def count_constant_components_2D(x, threshold = 1e-2, lower = -np.inf, upper = np.inf):
         """ Returns the number of piecewise constant components in a two-dimensional array
 
@@ -406,12 +438,19 @@ def count_constant_components_2D(x, threshold = 1e-2, lower = -np.inf, upper = n
         threshold : float
             Strict theshold on when the difference of neighbouring values is considered zero.
 
-        lower : float
+        lower : float or numpy.ndarray
             Piecewise constant components below this value are not counted.
 
-        upper : float
+        upper : float or numpy.ndarray
             Piecewise constant components above this value are not counted.
         """
+
+        if not isinstance(lower, np.ndarray):
+            lower = lower*np.ones_like(x)
+
+        if not isinstance(upper, np.ndarray):
+            upper = upper*np.ones_like(x)
+
         filled = np.zeros_like(x, dtype = int)
         counter = 0
 
@@ -438,8 +477,51 @@ def process(i, j):
         for i in range(x.shape[0]):
             for j in range(x.shape[1]):
                 if filled[i,j] == 0:
-                    if x[i,j] > lower and x[i,j] < upper:
+                    if x[i,j] > lower[i,j] and x[i,j] < upper[i,j]:
                         counter += 1
                     process(i, j)
         return counter
-                    
\ No newline at end of file
+                    
+
+
+def piecewise_linear_1D_DoF(x, threshold = 1e-5, exception_zero = False, exception_flat = False):
+        """ Returns the degrees of freedom of a piecewise linear signal.
+        Assuming linear interpolation, this corresponds to the number of non-differentiable points, including end-points.
+
+        Parameters
+        ----------
+        x : `np.ndarray` 
+            1D Array to compute degrees of freedom of.
+
+        threshold : float
+            Strict theshold on when values are considered zero.
+
+        exception_zero : Boolean
+            Whether a zero piecewise linear components should be considered.
+
+        exception_flat : Boolean
+            Whether a flat piecewise linear components should be considered.
+        """
+
+        differences = x[1:] - x[:-1]
+        double_differences = differences[1:] - differences[:-1]
+
+        joints = [True] + [np.abs(d) >= threshold for d in double_differences] + [True]
+        
+        if not exception_zero and not exception_flat:
+            return np.sum(joints)
+        elif exception_zero:
+            return np.sum(joints and [np.abs(v) < threshold for v in x])
+        elif exception_flat:
+            prev_joint = None
+            counter = 0
+            for i in range(len(joints)):
+                if joints[i]:
+                    counter += 1
+                    if prev_joint is None:
+                        prev_joint = i
+                    else:
+                        if np.abs(x[i] - x[prev_joint]) < threshold:
+                            counter -= 1
+                        prev_joint = i
+            return counter
diff --git a/demos/demo38_truncated_normal.py b/demos/demo38_truncated_normal.py
index 7f3b07219..58dce30fa 100644
--- a/demos/demo38_truncated_normal.py
+++ b/demos/demo38_truncated_normal.py
@@ -12,8 +12,8 @@
 p = cuqi.distribution.TruncatedNormal(
     mean=np.array([0, 0]),
     std=np.array([1, 1]),
-    low=np.array([0, -np.Inf]),
-    high=np.array([np.Inf, 0]))
+    low=np.array([0, -np.inf]),
+    high=np.array([np.inf, 0]))
 
 plt.figure()
 plot_2D_density(p, -5, 5, -5, 5)
@@ -40,8 +40,8 @@
 x = cuqi.distribution.TruncatedNormal(
     mean=np.array([0, 0]),
     std=np.array([1, 1]),
-    low=np.array([0, -np.Inf]),
-    high=np.array([np.Inf, 0]))
+    low=np.array([0, -np.inf]),
+    high=np.array([np.inf, 0]))
 # the data distribution
 b = cuqi.distribution.Gaussian(A@x, 0.1)
 
diff --git a/pyproject.toml b/pyproject.toml
index e7071d118..8156a7d1c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,7 @@ maintainers = [
     {name = "Chao Zhang", email = "chaz@dtu.dk"}
 ]
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 license = {file = "LICENSE"}
 dynamic = ["dependencies", "version"]
 
diff --git a/requirements.txt b/requirements.txt
index 3ae2538d6..540fe9b91 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
 matplotlib
 numpy>=1.17.0
-scipy<1.13
+scipy==1.12.0; python_version <= '3.9'
+scipy; python_version > '3.9'
 arviz
 tqdm


Test patch:
diff --git a/tests/test_bayesian_inversion.py b/tests/test_bayesian_inversion.py
index 97d27e854..4ea132186 100644
--- a/tests/test_bayesian_inversion.py
+++ b/tests/test_bayesian_inversion.py
@@ -2,7 +2,9 @@
 
 import pytest
 import numpy as np
+import scipy
 import sys
+from packaging import version
 
 from cuqi.testproblem import Deconvolution1D
 from cuqi.distribution import Gaussian, GMRF, CMRF, LMRF, Gamma
@@ -25,6 +27,14 @@ def test_TP_BayesianProblem_sample(copy_reference, TP_type, phantom, prior, Ns,
     # SKIP NUTS test if not windows (for now)
     if isinstance(prior, CMRF) and not sys.platform.startswith('win'):
         pytest.skip("NUTS(CMRF) regression test is not implemented for this platform")
+    elif (
+        isinstance(prior, RegularizedGaussian) 
+        or isinstance(prior, RegularizedGMRF)
+    ) and version.parse(scipy.__version__) >= version.parse('1.13.0'):
+        pytest.skip(
+            "Saved reference data is not consistent with scipy>=1.13.0 "
+            "(due to different default behaviour of scipy.optimize)"
+        )
 
     np.random.seed(19937)
 
diff --git a/tests/test_distribution.py b/tests/test_distribution.py
index 358d4a265..9a61265e1 100644
--- a/tests/test_distribution.py
+++ b/tests/test_distribution.py
@@ -203,7 +203,6 @@ def test_Uniform_logpdf(low, high, toeval, expected):
     UD = cuqi.distribution.Uniform(low, high)
     assert np.allclose(UD.logpdf(toeval), expected) 
 
-
 @pytest.mark.parametrize("low,high,expected",[(np.array([1, .5]), 
                                                np.array([2, 2.5]),
                                                np.array([[1.5507979 , 1.29090474, 1.89294695],
@@ -217,6 +216,19 @@ def test_Uniform_sample(low, high, expected):
     print(cuqi_samples)
     assert np.allclose(cuqi_samples.samples, expected) 
 
+def test_UnboundedUniform_logpdf_and_gradient():
+    x1 = cuqi.distribution.UnboundedUniform(geometry=1)
+    assert np.allclose(x1.logpdf(np.array([1])), 1.0)
+    assert np.allclose(x1.gradient(np.array([0])), np.array([0]))
+    x2 = cuqi.distribution.UnboundedUniform(geometry=cuqi.geometry.Discrete(2))
+    assert np.allclose(x2.logpdf(np.array([1, 2])), 1.0)
+    assert np.allclose(x2.gradient(np.array([0, 0])), np.array([0, 0]))
+
+def test_UnboundedUniform_sample():
+    x = cuqi.distribution.UnboundedUniform(geometry=2)
+    with pytest.raises(NotImplementedError, match="Cannot sample from UnboundedUniform distribution"):
+        x.sample()
+
 @pytest.mark.parametrize("distribution, kwargs",
                          [(cuqi.distribution.Uniform, 
                           {'low':np.array([2, 2.5, 3,5]),
@@ -665,7 +677,7 @@ def beta_likelihood():
         lambda x: x**2,
         range_geometry=1,
         domain_geometry=1,
-        gradient=lambda direction, wrt: 2*wrt*direction)
+        gradient=lambda direction, x: 2*x*direction)
     
     # set a gaussian prior
     x = cuqi.distribution.Gaussian(0, 1)
diff --git a/tests/test_distributions_shape.py b/tests/test_distributions_shape.py
index 7d0f69d1c..7a5b13780 100644
--- a/tests/test_distributions_shape.py
+++ b/tests/test_distributions_shape.py
@@ -13,6 +13,7 @@
     "Posterior",
     "Distribution",
     "JointGaussianSqrtPrec",
+    "UnboundedUniform",
 ]
 
 # Define cases to skip (these are TODO)
diff --git a/tests/test_implicit_priors.py b/tests/test_implicit_priors.py
index 4d4842e0f..2a8b4efa6 100644
--- a/tests/test_implicit_priors.py
+++ b/tests/test_implicit_priors.py
@@ -1,6 +1,8 @@
 import cuqi
 import numpy as np
+import scipy
 import pytest
+from packaging import version
 
 def test_RegularizedGaussian_default_init():
     """ Test that the implicit regularized Gaussian requires at least 1 regularization argument """
@@ -184,4 +186,68 @@ def test_RegularizedGaussian_double_preset():
     # Check whether the constructed proximal list is of correct size
     assert len(x.proximal) == 2
     assert len(x.proximal[0]) == 2
-    assert len(x.proximal[1]) == 2
\ No newline at end of file
+    assert len(x.proximal[1]) == 2
+
+@pytest.mark.skipif(version.parse(scipy.__version__) < version.parse('1.12.0'), reason="isotonic_regression in not available in scipy < 1.12.0")
+def test_regression_increasing():
+    """ Regression test for the increasing constraints in RegularizedGaussian"""
+
+    np.random.seed(1337)
+
+    n = 5
+    A, _, _ = cuqi.testproblem.Deconvolution1D(dim=n, phantom='square').get_components()
+
+    y_obs = np.array([0.1, 0.2, 0.3, 0.4, 0.5])
+
+    
+    d = cuqi.distribution.Gamma(1, 1e-4)
+    x = cuqi.implicitprior.RegularizedGaussian(mean = np.zeros(n), prec = lambda d : d,
+                            constraint = "increasing",
+                            geometry = A.domain_geometry,)
+    y = cuqi.distribution.Gaussian(A@x, prec = 1000)
+
+    joint = cuqi.distribution.JointDistribution(d, y, x)
+    posterior = joint(y=y_obs)
+   
+    sampling_strategy = {
+        'x': cuqi.experimental.mcmc.RegularizedLinearRTO(maxit=50, penalty_parameter=20, adaptive = False),
+        'd': cuqi.experimental.mcmc.Conjugate()
+                        }
+    sampler = cuqi.experimental.mcmc.HybridGibbs(posterior, sampling_strategy)
+
+    sampler.sample(10)
+    samples = sampler.get_samples()
+    
+    assert np.allclose(samples['x'].samples[:,-1], [0.07151248, 0.08253823, 0.3236245, 0.44846062, 0.61276932])
+  
+def test_regression_convex():
+    """ Regression test for the convexity constraints in RegularizedGaussian"""
+
+    np.random.seed(1337)
+
+    n = 5
+    A, _, _ = cuqi.testproblem.Deconvolution1D(dim=n, phantom='square').get_components()
+
+    y_obs = np.array([1, 0, -1, 2, 6])
+
+    
+    d = cuqi.distribution.Gamma(1, 1e-4)
+    x = cuqi.implicitprior.RegularizedGaussian(mean = np.zeros(n), prec = lambda d : d,
+                            constraint = "convex",
+                            geometry = A.domain_geometry,)
+    y = cuqi.distribution.Gaussian(A@x, prec = 1000)
+
+    joint = cuqi.distribution.JointDistribution(d, y, x)
+    posterior = joint(y=y_obs)
+   
+    sampling_strategy = {
+        'x': cuqi.experimental.mcmc.RegularizedLinearRTO(maxit=50, penalty_parameter=20, adaptive = False),
+        'd': cuqi.experimental.mcmc.Conjugate()
+                        }
+    sampler = cuqi.experimental.mcmc.HybridGibbs(posterior, sampling_strategy)
+
+    sampler.sample(10)
+    samples = sampler.get_samples()
+    
+    assert np.allclose(samples['x'].samples[:,-1], [-323.91409888, -194.50250442, -35.74426677, 167.01548949, 395.18988251])
+  
diff --git a/tests/test_model.py b/tests/test_model.py
index 0bb44beb4..41fc024bb 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -1,11 +1,1325 @@
 from cuqi.array import CUQIarray
-from cuqi.geometry import Continuous1D
+from cuqi.samples import Samples
+from cuqi.geometry import Continuous1D, Discrete
 import numpy as np
 import scipy as sp
 import cuqi
 import pytest
 from scipy import optimize
-from copy import copy
+from copy import copy, deepcopy
+from cuqi.geometry import _identity_geometries, _DefaultGeometry1D, _DefaultGeometry2D, Geometry, Discrete, Image2D
+from cuqi.utilities import force_ndarray
+from cuqi.experimental.geometry import _ProductGeometry
+
+def test_PDE_model_multiple_input():
+    """ Test that the PDE model can accept multiple inputs specified as positional arguments or keyword arguments """
+    pde_test_model = MultipleInputTestModel.helper_build_steady_state_PDE_test_model()
+    pde_test_model.populate_model_variations()
+    CUQI_pde = pde_test_model.model_variations[0] # PDE model with multiple inputs
+
+    # Check that the model has correct parameter name
+    assert CUQI_pde._non_default_args == ['mag', 'kappa_scale']
+
+    # Check that we can provide parameter names when evaluating the model
+    output1 = CUQI_pde(mag=2, kappa_scale=2)
+
+    # And check that we can provide positional arguments
+    output2 = CUQI_pde(2, 2)
+
+    # Check that the two outputs are the same
+    assert np.allclose(output1, output2)
+
+def test_constructing_gradient_from_jacobian():
+    """ Test that the gradient is correctly constructed from the
+    jacobian when only the jacobian is specified """
+
+    model = cuqi.model.Model(
+        lambda x, y: x * 2 + y,
+        domain_geometry=(Discrete(1), Discrete(1)),
+        range_geometry=1,
+        jacobian=lambda x, y: (np.array([2]), np.array([1])),
+    )
+
+    # Evaluate the gradient which was constructed from the jacobian:
+    grad1 = model.gradient(x=1, y=1, direction=np.array([2]))
+    # Stack grad1
+    grad1 = np.hstack([grad1[k] for k in grad1])
+
+    # Compute the gradient using the FD approximation:
+
+    # Slightly different form of the forward function
+    # that takes a single input (which approx_derivative requires)
+    forward2 = lambda x: x[0] * 2 + x[1]
+    grad2 = cuqi.utilities.approx_derivative(
+        forward2, np.array([1, 1]), np.array([2]), 1e-6
+    )
+
+    assert np.allclose(grad1, grad2)
+
+@pytest.mark.parametrize(
+    "obj_type", ['distribution', 'random_variable'])
+def test_model_updates_parameters_names_to_follow_distribution_or_random_variable_names(obj_type):
+    """Test that the model changes the parameter names if given distributions or random variables as input with new parameter names"""
+
+    def forward(x, y):
+        return y * x[0] + x[1]
+
+    def gradient_x(direction, x, y):
+        return direction * np.array([y, 1])
+
+    def gradient_y(direction, x, y):
+        return direction * x[0]
+
+    # forward model inputs
+    input1 = np.array([1, 2])
+    input2 = 1
+    direction = 3
+
+    model = cuqi.model.Model(
+        forward=forward,
+        range_geometry=1,
+        domain_geometry=(Continuous1D(2), Discrete(1)),
+        gradient=(gradient_x, gradient_y),
+    )
+
+    # assert model parameter names are 'x' and 'y'
+    assert model._non_default_args == ["x", "y"]
+    
+    if obj_type == 'distribution':
+        # create two random distributions a and b
+        a = cuqi.distribution.Gaussian(np.zeros(2), 1)
+        b = cuqi.distribution.Gaussian(np.zeros(1), 1)
+    elif obj_type == 'random_variable':
+        # create two random variables a and b
+        a = cuqi.distribution.Gaussian(np.zeros(2), 1).rv
+        b = cuqi.distribution.Gaussian(np.zeros(1), 1).rv
+    else:
+        raise ValueError("obj_type must be either 'distribution' or 'random_variable'")
+
+    model_a_b = model(a, b)
+
+    # assert model still has parameter names 'x' and 'y'
+    assert model._non_default_args == ["x", "y"]
+
+    # assert new model parameter names are 'a' and 'b'
+    assert model_a_b._non_default_args == ["a", "b"]
+
+    # assert the two models output are equal
+    model_output = model(x=input1, y=input2)
+    model_a_b_output = model_a_b(b=input2, a=input1)
+    assert np.allclose(model_output, model_a_b_output)
+
+    # assert the two models gradient are equal
+    model_grad_v1 = list(model.gradient(direction, x=input1, y=input2).values())
+    model_grad_v2 = list(model.gradient(direction, input1, input2).values())
+    model_grad_v3 = list(model.gradient(direction, y=input2, x=input1).values())
+    model_a_b_grad_v1 = list(model_a_b.gradient(direction, a=input1, b=input2).values())
+    model_a_b_grad_v2 = list(model_a_b.gradient(direction, input1, input2).values())
+    model_a_b_grad_v3 = list(model_a_b.gradient(direction, b=input2, a=input1).values())
+
+    # assert all gradients are equal
+    for i in range(len(model_grad_v1)):
+        assert np.allclose(model_grad_v1[i], model_grad_v2[i])
+        assert np.allclose(model_grad_v1[i], model_grad_v3[i])
+        assert np.allclose(model_grad_v1[i], model_a_b_grad_v1[i])
+        assert np.allclose(model_grad_v1[i], model_a_b_grad_v2[i])
+        assert np.allclose(model_grad_v1[i], model_a_b_grad_v3[i])
+
+    # compare to FD gradient
+    grad_FD = cuqi.utilities.approx_derivative(
+        lambda par: model.forward(par[:2], par[-1]),
+        np.hstack([input1, np.array([input2])]),
+        np.array([direction]),
+    )
+    assert np.allclose(model_grad_v1[0], grad_FD[:2])
+    assert np.allclose(model_grad_v1[1], grad_FD[-1])
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        cuqi.testproblem.Deconvolution1D().model,
+        cuqi.model.LinearModel(lambda x: x, lambda y: y, 1, 1),
+    ],
+)
+@pytest.mark.parametrize("obj_type", ['distribution', 'random_variable'])
+def test_linear_model_updates_parameters_names_to_follow_distribution_or_random_variable_names(
+    model, obj_type
+):
+    """Test that the linear model changes parameter names if given a distribution or a random variable as input with new parameter names"""
+
+    model_input = np.random.randn(model.domain_dim)
+
+    # assert model parameter names are 'x'
+    assert model._non_default_args == ["x"]
+    
+    if obj_type == 'distribution':
+        # create a random distribution a
+        a = cuqi.distribution.Gaussian(np.zeros(model.domain_dim), 1)
+    elif obj_type == 'random_variable':
+        # create a random variable a
+        a = cuqi.distribution.Gaussian(np.zeros(model.domain_dim), 1).rv
+    else:
+        raise ValueError("obj_type must be either 'distribution' or 'random_variable'")
+
+    model_a = model(a)
+
+    # assert model still has parameter names 'x'
+    assert model._non_default_args == ["x"]
+    # assert new model parameter names are 'a'
+    assert model_a._non_default_args == ["a"]
+
+    # assert the two models output are equal
+    assert np.allclose(model(x=model_input), model_a(a=model_input))
+    # assert _forward_func_no_shift also works with new parameter names
+    assert np.allclose(
+        model._forward_func_no_shift(x=model_input),
+        model_a._forward_func_no_shift(a=model_input),
+    )
+    # and also when inputs are positional arguments
+    assert np.allclose(
+        model._forward_func_no_shift(model_input),
+        model_a._forward_func_no_shift(model_input),
+    )
+
+    # assert the two models gradient are equal
+    direction = np.random.randn(model.range_dim)
+    model_grad = model.gradient(direction, x=model_input)
+    model_a_grad = model_a.gradient(direction, a=model_input)
+    assert np.allclose(model_grad, model_a_grad)
+
+    # assert the two models adjoint are equal
+    model_adjoint = model.adjoint(y=model_input)
+    model_a_adjoint = model_a.adjoint(y=model_input)
+    assert np.allclose(model_adjoint, model_a_adjoint)
+
+@pytest.mark.parametrize(
+    "gradient_with_incorrect_signature",
+    [
+        (lambda direction, x, y, z: direction),
+        (lambda direction, x: direction),
+        (lambda direction, y, x: direction),
+        (lambda direction, x, w: direction),
+    ],
+)
+def test_wrong_gradient_signature_raises_error_at_model_initialization(
+    gradient_with_incorrect_signature,
+):
+    """Test that an error is raised if the gradient signature is wrong"""
+
+    def forward(x, y):
+        return x
+
+    def gradient_with_correct_signature(direction, x, y):
+        return direction
+
+    cuqi.model.Model(
+        forward, 1, (Discrete(1), Discrete(1)), gradient=gradient_with_correct_signature
+    )
+
+    with pytest.raises(
+        ValueError,
+        match=r"Gradient function signature should be \['direction', 'x', 'y'\]",
+    ):
+        cuqi.model.Model(
+            forward,
+            1,
+            (Discrete(1), Discrete(1)),
+            gradient=gradient_with_incorrect_signature,
+        )
+
+def test_model_raises_error_when_domain_geometry_is_inconsistent_with_forward_signature_at_initialization():
+    """Test that the model raises an error if the domain geometry is inconsistent with the forward function signature at initialization"""
+
+    def forward(a, b):
+        return a * b
+
+    with pytest.raises(
+        ValueError,
+        match=r"The forward operator input is specified by more than one argument. This is only supported for domain geometry of type tuple",
+    ):
+        cuqi.model.Model(forward, 1, 2)
+
+def test_evaluating_model_at_distribution_with_non_unique_names_raises_error():
+    """Test that an error is raised if the model is evaluated at distributions with non-unique parameter names"""
+
+    def forward(a, b):
+        return a*b
+
+    model = cuqi.model.Model(forward, 1, (Discrete(1), Discrete(1)))
+
+    # Create distributions with non-unique parameter names
+    a = cuqi.distribution.Gaussian(0, 1, name="x")
+    b = cuqi.distribution.Gaussian(0, 1, name="x")
+
+    with pytest.raises(
+        ValueError,
+        match=r"Attempting to match parameter name of Model with given distributions, but distribution names are not unique. Please provide unique names for the distributions.",
+    ):
+        model(a, b)
+
+def test_evaluating_model_at_random_variables_with_non_unique_names_raises_error():
+    """Test that an error is raised if the model is evaluated at random variables with non-unique parameter names"""
+
+    def forward(a, b):
+        return a*b
+
+    model = cuqi.model.Model(forward, 1, (Discrete(1), Discrete(1)))
+
+    # Create distributions with non-unique parameter names
+    a = cuqi.distribution.Gaussian(0, 1, name="x").rv
+    b = cuqi.distribution.Gaussian(0, 1, name="x").rv
+
+    with pytest.raises(
+        ValueError,
+        match=r"Attempting to match parameter name of Model with given random variables, but random variables names are not unique. Please provide unique names for the random variables.",
+    ):
+        model(a, b)
+
+@pytest.mark.parametrize(
+    "x, y, z, forward, domain_geometry, range_geometry, expected_value, expected_type",
+    [
+        (
+            np.array([1, 3, 4]),
+            np.array([1, 3]),
+            3,
+            lambda x, y, z: x * y[0] + z * y[1],
+            cuqi.experimental.geometry._ProductGeometry(
+                cuqi.geometry.Continuous1D(3),
+                cuqi.geometry.Continuous1D(2),
+                cuqi.geometry.Continuous1D(3),
+            ),
+            cuqi.geometry.Continuous1D(3),
+            np.array([10, 12, 13]),
+            np.ndarray,
+        )
+    ],
+)
+class TestForwardWithMultipleInputs:
+    def test_forward_with_multiple_inputs(
+        self,
+        x,
+        y,
+        z,
+        forward,
+        domain_geometry,
+        range_geometry,
+        expected_value,
+        expected_type,
+    ):
+        """Test that the forward method can handle multiple inputs and return
+        the correct output type (even if order of kwargs inputs is switched)"""
+
+        model = cuqi.model.Model(
+            forward=forward,
+            domain_geometry=domain_geometry,
+            range_geometry=range_geometry,
+        )
+
+        fwd = model.forward(x, y, z)
+        fwd2 = model(y=y, x=x, z=z)
+        fwd3 = model.forward(x, y, z, is_par=True)
+        fwd4 = model(z=z, x=x, y=y, is_par=True)
+        assert isinstance(fwd, expected_type)
+        assert np.allclose(fwd, expected_value)
+        assert np.allclose(fwd, fwd2)
+        assert np.allclose(fwd, fwd3)
+        assert np.allclose(fwd, fwd4)
+
+    def test_forward_with_multiple_inputs_error_when_mixing_args_and_kwargs(
+        self,
+        x,
+        y,
+        z,
+        forward,
+        domain_geometry,
+        range_geometry,
+        expected_value,
+        expected_type,
+    ):
+        """Test that the forward method raises an error when mixing positional
+        arguments and keyword arguments"""
+
+        model = cuqi.model.Model(
+            forward=forward,
+            domain_geometry=domain_geometry,
+            range_geometry=range_geometry,
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=r"The model input is specified both as positional and keyword arguments. This is not supported",
+        ):
+            model.forward(x, y=y, z=z)
+
+class MultipleInputTestModel:
+    """Class representing a test model with ingredients to set up variations of the model. For example, model with gradient, or with jacobian, etc. All the variations shares the same forward map, the domain and range geometry."""
+
+    def __init__(self):
+        self.model_class = None  # the cuqi.model class to set up the model
+        self.forward_map = None
+        self.pde = None  # in case of PDE model
+        self.gradient_form1 = None  # callable
+        self.gradient_form2 = None  # tuple of callables
+        self.gradient_form2_incomplete = (
+            None  # tuple of callables with some None elements
+        )
+        self.jacobian_form1 = None  # callable
+        self.jacobian_form2 = None  # tuple of callables
+        self.jacobian_form2_incomplete = (
+            None  # tuple of callables with some None elements
+        )
+        self.domain_geometry = None
+        self.range_geometry = None
+        self.test_data = []  # list of TestCase objects holding the test data
+        self.model_variations = (
+            []
+        )  # list of model variations which all share the same forward map
+        self.input_bounds = None  # bounds for the test input values
+
+    @property
+    def has_gradient(self):
+        """Check if the MultipleInputTestModel object has a way of determining 
+        the model gradient. That is, if this object has gradient or jacobian
+        (which can be used to specify the gradient) function of any form"""
+        return (
+            self.gradient_form1 is not None
+            or self.gradient_form2 is not None
+            or self.gradient_form2_incomplete is not None
+            or self.jacobian_form1 is not None
+            or self.jacobian_form2 is not None
+            or self.jacobian_form2_incomplete is not None
+        )
+
+    def populate_model_variations(self):
+        """Populate the `model_variations` list with different variations of the model that share the same forward map but differ in some other aspect like gradient, jacobian, etc."""
+        if self.pde is not None:
+            self.populate_pde_model_variations()
+        elif self.forward_map is not None:
+            self.populate_general_model_variations()
+
+    def populate_general_model_variations(self):
+        """Populate the `model_variations` for general models that are not PDE models."""
+        # Model with forward only
+        model = self.model_class(
+            forward=self.forward_map,
+            domain_geometry=self.domain_geometry,
+            range_geometry=self.range_geometry,
+        )
+        model._do_test_gradient = False  # do not test this model for gradient
+        self.model_variations.append(model)
+
+        # Model with gradient of from 1 (callable)
+        if self.gradient_form1 is not None:
+            model = self.model_class(
+                forward=self.forward_map,
+                gradient=self.gradient_form1,
+                domain_geometry=self.domain_geometry,
+                range_geometry=self.range_geometry,
+            )
+            self.model_variations.append(model)
+
+        # Model with gradient of from 2 (tuple of callables)
+        if self.gradient_form2 is not None:
+            model = self.model_class(
+                forward=self.forward_map,
+                gradient=self.gradient_form2,
+                domain_geometry=self.domain_geometry,
+                range_geometry=self.range_geometry,
+            )
+            self.model_variations.append(model)
+
+        # Model with gradient of from 2 incomplete (tuple of callables with some None elements)
+        if self.gradient_form2_incomplete is not None:
+            model = self.model_class(
+                forward=self.forward_map,
+                gradient=self.gradient_form2_incomplete,
+                domain_geometry=self.domain_geometry,
+                range_geometry=self.range_geometry,
+            )
+            self.model_variations.append(model)
+
+        # Model with jacobian of from 1 (callable)
+        if self.jacobian_form1 is not None:
+            model = self.model_class(
+                forward=self.forward_map,
+                jacobian=self.jacobian_form1,
+                domain_geometry=self.domain_geometry,
+                range_geometry=self.range_geometry,
+            )
+            self.model_variations.append(model)
+
+        # Model with jacobian of from 2 (tuple of callables)
+        if self.jacobian_form2 is not None:
+            model = self.model_class(
+                forward=self.forward_map,
+                jacobian=self.jacobian_form2,
+                domain_geometry=self.domain_geometry,
+                range_geometry=self.range_geometry,
+            )
+            self.model_variations.append(model)
+
+        # Model with jacobian of from 2 incomplete (tuple of callables with some None elements)
+        if self.jacobian_form2_incomplete is not None:
+            model = self.model_class(
+                forward=self.forward_map,
+                jacobian=self.jacobian_form2_incomplete,
+                domain_geometry=self.domain_geometry,
+                range_geometry=self.range_geometry,
+            )
+            self.model_variations.append(model)
+
+    def populate_pde_model_variations(self):
+        """Populate the `model_variations` for PDE models."""
+        # Model with PDE (no gradient)
+        model = self.model_class(
+            self.pde,
+            domain_geometry=self.domain_geometry,
+            range_geometry=self.range_geometry,
+        )
+        model._do_test_gradient = False  # do not test this model for gradient
+        self.model_variations.append(model)
+
+    @staticmethod
+    def create_model_test_case_combinations():
+        """Create all combinations of test model variations and test cases
+        (test data) for the test models."""
+        model_test_case_combinations = []
+        test_model_list = []
+
+        # Model 1
+        test_model = MultipleInputTestModel.helper_build_three_input_test_model()
+        test_model.populate_model_variations()
+        TestCase.create_test_cases_for_test_model(test_model)
+        test_model_list.append(test_model)
+
+        # Model 2
+        test_model = MultipleInputTestModel.helper_build_steady_state_PDE_test_model()
+        test_model.populate_model_variations()
+        TestCase.create_test_cases_for_test_model(test_model)
+        test_model_list.append(test_model)
+
+        # Model 3
+        test_model = (
+            MultipleInputTestModel.helper_build_two_input_test_with_mapped_parameter_model()
+        )
+        test_model.populate_model_variations()
+        test_model.input_bounds = [
+            0.1,
+            0.9,
+        ]  # choose input from uniform distribution in [0.1, 0.9]
+        TestCase.create_test_cases_for_test_model(test_model)
+        test_model_list.append(test_model)
+
+        # Append all combinations of test model variations and test cases
+        # to model_test_case_combinations
+        for test_model in test_model_list:
+            for test_model_variation in test_model.model_variations:
+                for test_data_item in test_model.test_data:
+                    model_test_case_combinations.append(
+                        (test_model_variation, test_data_item)
+                    )
+
+        return model_test_case_combinations
+
+    @staticmethod
+    def helper_build_three_input_test_model():
+        """Build a MultipleInputTestModel with three inputs: x, y, z, and gradient and jacobian functions."""
+        test_model = MultipleInputTestModel()
+        test_model.forward_map = lambda x, y, z: x * y[0] + z * y[1]
+
+        # gradient with respect to x
+        def gradient_x(direction, x, y, z):
+            return direction * y[0]
+
+        # gradient with respect to y
+        def gradient_y(direction, x, y, z):
+            return np.array([direction @ x, direction @ z])
+
+        # gradient with respect to z
+        def gradient_z(direction, x, y, z):
+            return direction * y[1]
+
+        # gradient with respect to all inputs (form 1, callable)
+        def gradient_form1(direction, x, y, z):
+            grad_x = gradient_x(direction, x, y, z)
+            grad_y = gradient_y(direction, x, y, z)
+            grad_z = gradient_z(direction, x, y, z)
+            return (grad_x, grad_y, grad_z)
+
+        # Assign the gradient functions to the test model
+        test_model.gradient_form1 = gradient_form1
+        test_model.gradient_form2 = (gradient_x, gradient_y, gradient_z)
+        test_model.gradient_form2_incomplete = (gradient_x, None, gradient_z)
+
+        # jacobian with respect to x
+        def jacobian_x(x, y, z):
+            ones = np.ones_like(x)
+            return np.diag(y[0] * ones)
+
+        # jacobian with respect to y
+        def jacobian_y(x, y, z):
+            return np.array([x, z]).T
+
+        # jacobian with respect to z
+        def jacobian_z(x, y, z):
+            ones = np.ones_like(x)
+            return np.diag(y[1] * ones)
+
+        # jacobian with respect to all inputs (form 1, callable)
+        def jacobian_form1(x, y, z):
+            jac_x = jacobian_x(x, y, z)
+            jac_y = jacobian_y(x, y, z)
+            jac_z = jacobian_z(x, y, z)
+            return (jac_x, jac_y, jac_z)
+
+        # Assign the jacobian functions to the test model
+        test_model.jacobian_form1 = jacobian_form1
+        test_model.jacobian_form2 = (jacobian_x, jacobian_y, jacobian_z)
+        test_model.jacobian_form2_incomplete = (None, jacobian_y, jacobian_z)
+
+        test_model.domain_geometry = (
+            cuqi.geometry.Continuous1D(3),
+            cuqi.geometry.Continuous1D(2),
+            cuqi.geometry.Continuous1D(3),
+        )
+        test_model.range_geometry = cuqi.geometry.Continuous1D(3)
+
+        test_model.model_class = cuqi.model.Model
+        return test_model
+
+    @staticmethod
+    def helper_build_steady_state_PDE_test_model():
+        """Build a PDE model with a steady state Poisson equation and two inputs: mag and kappa_scale. This model does not have gradient or jacobian functions."""
+
+        # Poisson equation setup
+        dim = 20  # Number of nodes
+        L = 20  # Length of the domain
+        dx = L / (dim - 1)  # grid spacing
+        grid_sol = np.linspace(dx, L, dim - 1, endpoint=False)  # solution grid
+        grid_obs = grid_sol[5:]  # observation grid
+        source = lambda mag: mag * np.sin(
+            grid_sol
+        )  # source term depending on magnitude `mag`
+        kappa = np.ones(dim)  # kappa is the diffusivity
+
+        # Build the differential operator (depending on kappa_scale, a scale of the diffusivity)
+        FOFD_operator = (
+            cuqi.operator.FirstOrderFiniteDifference(dim - 1, bc_type="zero", dx=dx)
+            .get_matrix()
+            .todense()
+        )
+        diff_operator = (
+            lambda kappa_scale: FOFD_operator.T
+            @ np.diag(kappa_scale * kappa)
+            @ FOFD_operator
+        )
+
+        # Build the PDE form
+        poisson_form = lambda mag, kappa_scale: (
+            diff_operator(kappa_scale),
+            source(mag),
+        )
+
+        # Build the PDE object
+        CUQI_pde = cuqi.pde.SteadyStateLinearPDE(
+            poisson_form,
+            grid_sol=grid_sol,
+            grid_obs=grid_obs,
+            observation_map=lambda u: u**2,
+        )
+
+        # Build the test model
+        test_model = MultipleInputTestModel()
+        test_model.model_class = cuqi.model.PDEModel
+        test_model.pde = CUQI_pde
+        test_model.domain_geometry = (Discrete(["mag"]), Discrete(["kappa_scale"]))
+        test_model.range_geometry = Continuous1D(len(grid_obs))
+
+        return test_model
+
+    @staticmethod
+    def helper_build_two_input_test_with_mapped_parameter_model():
+        """Build a model with two inputs in which the inputs are mapped
+        via geometry mapping"""
+        test_model = MultipleInputTestModel()
+        test_model.forward_map = lambda a, b: np.array(
+            [a[0] ** 2 * b[0] * b[1], b[1] * a[0] * a[1]]
+        )
+
+        # gradient with respect to a
+        def gradient_a(direction, a, b):
+            return np.array(
+                [
+                    2 * a[0] * b[0] * b[1] * direction[0] + b[1] * a[1] * direction[1],
+                    b[1] * a[0] * direction[1],
+                ]
+            )
+
+        # gradient with respect to b
+        def gradient_b(direction, a, b):
+            return np.array(
+                [
+                    a[0] ** 2 * b[1] * direction[0],
+                    a[0] ** 2 * b[0] * direction[0] + a[0] * a[1] * direction[1],
+                ]
+            )
+
+        # gradient with respect to all inputs (form 1, callable)
+        def gradient_form1(direction, a, b):
+            grad_a = gradient_a(direction, a, b)
+            grad_b = gradient_b(direction, a, b)
+            return (grad_a, grad_b)
+
+        test_model.gradient_form1 = gradient_form1
+        test_model.gradient_form2 = (gradient_a, gradient_b)
+        test_model.gradient_form2_incomplete = (gradient_a, None)
+
+        test_model.jacobian_form1 = None
+        test_model.jacobian_form2 = None
+        test_model.jacobian_form2_incomplete = None
+
+        # Setting up domain and range geometries
+        geom_a = cuqi.geometry.MappedGeometry(
+            cuqi.geometry.Continuous1D(2),
+            map=lambda x: np.sin(x),
+            imap=lambda x: np.arcsin(x),
+        )
+        geom_a.gradient = lambda direction, x: direction @ np.diag(np.cos(x))
+
+        geom_b = cuqi.geometry.MappedGeometry(
+            cuqi.geometry.Continuous1D(2),
+            map=lambda x: x**3,
+            imap=lambda x: x ** (1 / 3),
+        )
+        geom_b.gradient = lambda direction, x: 3 * direction @ np.diag(x**2)
+        test_model.domain_geometry = (geom_a, geom_b)
+        test_model.range_geometry = cuqi.geometry.Continuous1D(2)
+
+        # Assign the model class
+        test_model.model_class = cuqi.model.Model
+        return test_model
+
+
+class TestCase:
+    """Class representing a test case for a test model. A test case consists of the input values, the expected output values, and the expected output types and error messages."""
+
+    def __init__(self, test_model):
+        self._test_model = test_model
+
+        # Non default arguments of the cuqi.model.Model object
+        self._non_default_args = self._test_model.model_variations[0]._non_default_args
+        self.forward_input = None
+        self.forward_input_stacked = None
+        self.direction = None
+        self.expected_fwd_output = None
+        self.expected_grad_output_value = None
+        self.expected_staked_grad_output_value = None
+        self.expected_fwd_output_type = None
+        self.expected_grad_output_type = None
+        self.FD_grad_output = None  # finite difference gradient for verification
+
+    @property
+    def model_forward_map(self):
+        """Returns the underlying forward map of the test model (the one that maps funvals to funvals). Note that for cuqi.model.Model objects, the underlying forward map, stored in the `_forward_func` attribute, maps funvals to funvals, while the `forward` method that wraps the underlying `_forward_func` maps parameter values to parameter values."""
+
+        return self._test_model.model_variations[0]._forward_func
+
+    def create_input(self):
+        """Create random input values for the test case, stored in two formats, a kwargs dictionary and a stacked numpy array."""
+        input_dict = {}  # kwargs dictionary
+
+        for i, arg in enumerate(self._non_default_args):
+            dim = self._test_model.domain_geometry[i].par_dim
+            # If case: input bounds are specified, sample from uniform distribution
+            if self._test_model.input_bounds is not None:
+                input_dict[arg] = np.random.uniform(
+                    [self._test_model.input_bounds[0]] * dim,
+                    [self._test_model.input_bounds[1]] * dim,
+                )
+            # Else case: sample from standard normal distribution
+            else:
+                input_dict[arg] = np.random.randn(dim)
+
+        self.forward_input = input_dict
+        self.forward_input_stacked = np.hstack([v for v in list(input_dict.values())])
+
+    def create_direction(self):
+        """Create a random direction for the gradient computation."""
+        self.direction = np.random.randn(self._test_model.range_geometry.par_dim)
+
+    def compute_expected_fwd_output(self):
+        """Compute the expected output of the forward map for the test case. Note that this does not use the method `forward` of the `cuqi.model.Model` object, but the underlying forward map stored in the `_forward_func` attribute which is the function supplied by the user (in Model class case) or built from the PDE object (in PDEModel case)."""
+        mapped_input = self._compute_mapped_input(
+            self.forward_input, self._test_model.domain_geometry
+        )
+        self.expected_fwd_output = self.model_forward_map(**mapped_input)
+
+    @staticmethod
+    def _compute_mapped_input(forward_input, domain_geom):
+        """Map the input values to function values using the par2fun method of the domain geometry."""
+        mapped_input = {}
+        for i, (k, v) in enumerate(forward_input.items()):
+            if hasattr(domain_geom[i], "map"):
+                mapped_input[k] = domain_geom[i].par2fun(v)
+            elif type(domain_geom[i]) in _identity_geometries:
+                mapped_input[k] = v
+            else:
+                raise NotImplementedError
+        return mapped_input
+
+    def compute_expected_grad_output(self):
+        """Compute the expected output of the gradient computation for the test case. Note that this does not use the method `gradient` of the `cuqi.model.Model` object, but the underlying gradient functions supplied by the user."""
+
+        domain_geom = self._test_model.domain_geometry
+
+        # Map input values to function values
+        forward_input_fun = deepcopy(self.forward_input)
+        forward_input_fun = self._compute_mapped_input(self.forward_input, domain_geom)
+
+        # Compute the expected gradient output
+        self.expected_grad_output_value = []
+        for i, (k, v) in enumerate(self.forward_input.items()):
+            self.expected_grad_output_value.append(
+                self._test_model.gradient_form2[i](self.direction, **forward_input_fun)
+            )
+            # If case: apply the chain rule for the gradient computation if the geometry has a gradient method
+            if hasattr(domain_geom[i], "gradient"):
+                self.expected_grad_output_value[-1] = domain_geom[i].gradient(
+                    self.expected_grad_output_value[-1], v
+                )
+            elif type(domain_geom[i]) not in _identity_geometries:
+                raise NotImplementedError
+        self.expected_grad_output_value = tuple(self.expected_grad_output_value)
+        self._compute_expected_stacked_grad_output()
+
+    def _compute_expected_stacked_grad_output(self):
+        """Compute the expected stacked gradient output for the test case. This is used for testing the gradient computation when the model flag _gradient_output_stacked is enabled."""
+        expected_staked_grad_output = np.hstack(
+            [
+                (
+                    v.to_numpy()
+                    if isinstance(v, CUQIarray)
+                    else force_ndarray(v, flatten=True)
+                )
+                for v in list(self.expected_grad_output_value)
+            ]
+        )
+        self.expected_staked_grad_output_value = expected_staked_grad_output
+
+    def compute_FD_grad_output(self):
+        """Compute the finite difference gradient for the test case, for verification."""
+        FD_grad_list = []
+        model = self._test_model.model_variations[0]
+        for k, v in self.forward_input.items():
+            # Create forward lambda function with one input x, representing
+            # the value of k to be used in computing the finite difference
+            # gradient with respect to k
+            forward_input = self.forward_input.copy()
+            del forward_input[k]
+            fwd = lambda x: model(**forward_input, **{k: x})
+
+            if isinstance(v, cuqi.array.CUQIarray):
+                v = v.to_numpy()
+
+            direction = self.direction
+            if isinstance(self.direction, cuqi.array.CUQIarray):
+                direction = self.direction.to_numpy()
+
+            FD_grad = cuqi.utilities.approx_derivative(fwd, v, direction)
+            FD_grad_list.append(FD_grad)
+
+        self.FD_grad_output = tuple(FD_grad_list)
+
+    @staticmethod
+    def create_test_cases_for_test_model(test_model):
+        # Case 1: all inputs are numpy arrays
+        test_case = TestCase(test_model)
+        test_case.create_input()
+        test_case.compute_expected_fwd_output()
+        test_case.expected_fwd_output_type = np.ndarray
+
+        if test_model.has_gradient:
+            test_case.create_direction()
+            test_case.compute_expected_grad_output()
+            test_case.compute_FD_grad_output()
+            test_case.expected_grad_output_type = np.ndarray
+
+        test_model.test_data.append(test_case)
+
+        # Case 2: all inputs are CUQIarrays
+        test_case = TestCase(test_model)
+        test_case.create_input()
+        for i, (k, v) in enumerate(test_case.forward_input.items()):
+            test_case.forward_input[k] = cuqi.array.CUQIarray(
+                v, geometry=test_model.domain_geometry[i]
+            )
+        test_case.compute_expected_fwd_output()
+        test_case.expected_fwd_output = cuqi.array.CUQIarray(
+            test_case.expected_fwd_output, geometry=test_model.range_geometry
+        )
+        test_case.expected_fwd_output_type = cuqi.array.CUQIarray
+
+        if test_model.has_gradient:
+            test_case.create_direction()
+            test_case.direction = cuqi.array.CUQIarray(
+                test_case.direction, geometry=test_model.range_geometry
+            )
+            test_case.compute_expected_grad_output()
+            test_case.compute_FD_grad_output()
+            test_case.expected_grad_output_type = cuqi.array.CUQIarray
+
+        test_model.test_data.append(test_case)
+
+        # Case 3: inputs are mix of CUQIarrays and numpy arrays
+        test_case = TestCase(test_model)
+        test_case.create_input()
+        for i, (k, v) in enumerate(test_case.forward_input.items()):
+            if i == 0:
+                test_case.forward_input[k] = cuqi.array.CUQIarray(
+                    v, geometry=test_model.domain_geometry[i]
+                )
+        test_case.compute_expected_fwd_output()
+        test_case.expected_fwd_output = cuqi.array.CUQIarray(
+            test_case.expected_fwd_output, geometry=test_model.range_geometry
+        )
+        test_case.expected_fwd_output_type = cuqi.array.CUQIarray
+
+        if test_model.has_gradient:
+            test_case.create_direction()
+            test_case.compute_expected_grad_output()
+            test_case.compute_FD_grad_output()
+            test_case.expected_grad_output_type = cuqi.array.CUQIarray
+
+        test_model.test_data.append(test_case)
+
+        # Case 4: same as previous case but direction is a CUQIarray
+        test_case = deepcopy(test_case)
+        if test_model.has_gradient:
+            test_case.direction = cuqi.array.CUQIarray(
+                test_case.direction, geometry=test_model.range_geometry
+            )
+
+        test_model.test_data.append(test_case)
+
+        # Case 5: inputs are mix of CUQIarrays and samples (should raise an error)
+        test_case = TestCase(test_model)
+        test_case.create_input()
+        for i, (k, v) in enumerate(test_case.forward_input.items()):
+            if i == 0:
+                v2 = 1.3 * v
+                samples = np.vstack([v, v2]).T
+                test_case.forward_input[k] = cuqi.samples.Samples(
+                    samples, geometry=test_model.domain_geometry[i]
+                )
+            else:
+                test_case.forward_input[k] = cuqi.array.CUQIarray(
+                    v, geometry=test_model.domain_geometry[i]
+                )
+        test_case.expected_fwd_output = TypeError(
+            "If applying the function to Samples, all inputs should be Samples."
+        )
+
+        if test_model.has_gradient:
+            test_case.create_direction()
+            test_case.expected_grad_output_value = NotImplementedError(
+                "Gradient is not implemented for input of type Samples."
+            )
+
+        test_model.test_data.append(test_case)
+
+        # Case 6: inputs are samples (should work for forward but not for gradient)
+        test_case = TestCase(test_model)
+
+        # Create model input as samples
+        test_case.create_input()
+        for i, (k, v) in enumerate(test_case.forward_input.items()):
+            v2 = 1.5 * v
+            samples = np.vstack([v, v2]).T
+            test_case.forward_input[k] = cuqi.samples.Samples(
+                samples, geometry=test_model.domain_geometry[i]
+            )
+        # Create stacked samples input (not supported for forward or gradient)
+        # We use it in testing that it raises an error
+        stacked_input_samples = []
+        for i in range(2):
+            stacked_input_samples.append(
+                np.hstack([v.samples[:, i] for v in test_case.forward_input.values()])
+            )
+        test_case.forward_input_stacked = cuqi.samples.Samples(
+            np.vstack([stacked_input_samples]).T, geometry=test_model.domain_geometry
+        )
+
+        # Compute expected forward output
+        expected_fwd_output = []
+        for i in range(2):
+            input_i = {k: v.samples[:, i] for k, v in test_case.forward_input.items()}
+            mapped_input_i = TestCase._compute_mapped_input(
+                input_i, test_model.domain_geometry
+            )
+            expected_fwd_output.append(test_case.model_forward_map(**mapped_input_i))
+        expected_fwd_output = np.vstack(expected_fwd_output).T
+        test_case.expected_fwd_output = cuqi.samples.Samples(
+            expected_fwd_output, geometry=test_model.range_geometry
+        )
+        test_case.expected_fwd_output_type = cuqi.samples.Samples
+
+        if test_model.has_gradient:
+            test_case.create_direction()
+            test_case.expected_grad_output_value = NotImplementedError(
+                "Gradient is not implemented for input of type Samples."
+            )
+
+        test_model.test_data.append(test_case)
+
+        # Case 7: inputs are samples but of different length (should raise an error)
+        test_case = TestCase(test_model)
+
+        # Create model input as samples
+        test_case.create_input()
+        for i, (k, v) in enumerate(test_case.forward_input.items()):
+            if i == 0:  # Case i=0: two samples
+                v2 = 1.5 * v
+                samples = np.vstack([v, v2]).T
+                test_case.forward_input[k] = cuqi.samples.Samples(
+                    samples, geometry=test_model.domain_geometry[i]
+                )
+            else:  # Case i>0: three samples
+                v2 = 1.5 * v
+                v3 = 2 * v
+                samples = np.vstack([v, v2, v3]).T
+                test_case.forward_input[k] = cuqi.samples.Samples(
+                    samples, geometry=test_model.domain_geometry[i]
+                )
+
+        test_case.expected_fwd_output = ValueError(
+            "If applying the function to Samples, all inputs should have the same number of samples."
+        )
+
+        if test_model.has_gradient:
+            test_case.create_direction()
+            test_case.expected_grad_output_value = NotImplementedError(
+                "Gradient is not implemented for input of type Samples."
+            )
+
+        test_model.test_data.append(test_case)
+
+
+# Create all combinations of test model variations and test cases and
+# store them in model_test_case_combinations to be used in the tests
+model_test_case_combinations = (
+    MultipleInputTestModel.create_model_test_case_combinations()
+)
+
+
+def helper_function_for_printing_test_cases(model_test_case_combinations):
+    """Helper function to print the test cases for debugging. This function
+    is not used in the tests, but can be used by developers to print the
+    test cases for debugging."""
+    for i, (test_model, test_data) in enumerate(model_test_case_combinations):
+        print("### Test Case {} ###".format(i))
+        print("## test_model ##")
+        print(test_model)
+        print("## test_model._forward_func ##")
+        print(test_model._forward_func)
+        print("## test_model._gradient_func ##")
+        print(test_model._gradient_func)
+        print("## test_model._domain_geometry ##")
+        print(test_model._domain_geometry)
+        print("## test_model._range_geometry ##")
+        print(test_model._range_geometry)
+        print("=" * 50)
+        print("## test_data ##")
+        print(test_data)
+        print("## test_data.forward_input ##")
+        print(test_data.forward_input)
+        print("## test_data.expected_fwd_output_type ##")
+        print(test_data.expected_fwd_output_type)
+        print("## test_data.expected_fwd_output ##")
+        print(test_data.expected_fwd_output)
+        print("## test_data.expected_grad_output_type ##")
+        print(test_data.expected_grad_output_type)
+        print("## test_data.expected_grad_output_value ##")
+        print(test_data.expected_grad_output_value)
+        print("#" * 50)
+        print("#" * 50)
+
+
+@pytest.mark.parametrize("test_model, test_data", model_test_case_combinations)
+def test_forward_of_multiple_input_model_is_correct(test_model, test_data):
+    """Test that the forward method can handle multiple inputs when evaluated on the test data and return the correct output type and value"""
+    assert isinstance(test_model, cuqi.model.Model)
+    assert isinstance(test_data, TestCase)
+
+    # If case: cases where we can apply the forward function without raising an error
+    if not isinstance(
+        test_data.expected_fwd_output, (NotImplementedError, TypeError, ValueError)
+    ):
+        fwd_output = test_model(**test_data.forward_input)
+
+        # Assert output has the expected type
+        assert isinstance(fwd_output, test_data.expected_fwd_output_type)
+
+        # Assert output has the expected value
+        if isinstance(fwd_output, np.ndarray):
+            assert np.allclose(fwd_output, test_data.expected_fwd_output)
+        elif isinstance(fwd_output, cuqi.samples.Samples):
+            assert np.allclose(
+                fwd_output.samples, test_data.expected_fwd_output.samples
+            )
+        else:
+            raise NotImplementedError(
+                "Checks for other types of outputs not implemented."
+            )
+
+    # Else case: cases where applying the forward function raises an error
+    else:
+        with pytest.raises(
+            type(test_data.expected_fwd_output),
+            match=str(test_data.expected_fwd_output),
+        ):
+            fwd_output = test_model(**test_data.forward_input)
+
+
+# Create a sublist of model_test_case_combinations for cases where the expected
+# forward output is not an error
+model_test_case_combinations_no_forward_error = [
+    (test_model, test_data)
+    for test_model, test_data in model_test_case_combinations
+    if not isinstance(
+        test_data.expected_fwd_output, (NotImplementedError, TypeError, ValueError)
+    )
+]
+
+
+@pytest.mark.parametrize(
+    "test_model, test_data", model_test_case_combinations_no_forward_error
+)
+def test_forward_of_multiple_input_model_is_correct_when_input_is_stacked(
+    test_model, test_data
+):
+    """Test that the forward method can handle stacked inputs and return the correct
+    output type and value"""
+    assert isinstance(test_model, cuqi.model.Model)
+    assert isinstance(test_data, TestCase)
+
+    # If case: the output is expected to be a numpy array
+    if isinstance(test_data.expected_fwd_output, np.ndarray):
+        assert len(test_data.forward_input_stacked) == test_model.domain_dim
+        fwd_output_stacked_inputs = test_model.forward(test_data.forward_input_stacked)
+        assert isinstance(fwd_output_stacked_inputs, np.ndarray)
+        assert np.allclose(fwd_output_stacked_inputs, test_data.expected_fwd_output)
+    # Else case: case where the input (and hence expected output) are Samples object.
+    # For this case, the forward method should raise an error since applying the
+    # forward function to a Samples object of stacked input is not supported.
+    elif isinstance(test_data.expected_fwd_output, Samples):
+        with pytest.raises(
+            ValueError,
+            match=r"The model input is specified by a Samples object that cannot be split into multiple arguments corresponding to the non_default_args",
+        ):
+            test_model.forward(test_data.forward_input_stacked)
+    else:
+        raise NotImplementedError
+
+    # Assert evaluating forward on stacked input with wrong dimension raises error
+    if isinstance(test_data.expected_fwd_output, np.ndarray):
+        with pytest.raises(
+            ValueError,
+            match=r"The number of positional arguments does not match the number of non-default arguments of the model. Additionally, the model input is specified by a single argument that cannot be split into multiple arguments matching the expected non_default_args",
+        ):
+            test_model.forward(test_data.forward_input_stacked[:-1])
+
+
+@pytest.mark.parametrize(
+    "test_model, test_data", model_test_case_combinations_no_forward_error
+)
+def test_forward_of_multiple_input_model_applied_to_funvals_input_is_correct(
+    test_model, test_data
+):
+    """Test that the forward method can handle multiple inputs some of which are
+    function values and return the correct output"""
+
+    # par input
+    par_input = test_data.forward_input
+    model_output_par_input = test_model(**par_input)
+
+    # fun input (note because some of the test data cases have a mix of numpy
+    # arrays and cuqi arrays, only the cuqi arrays are converted to funvals
+    # below, hence it will be a case of mixed input of funvals and par values)
+    fun_input = {
+        k: v.funvals if (isinstance(v, CUQIarray) or isinstance(v, Samples)) else v
+        for k, v in par_input.items()
+    }
+    model_output_fun_input = test_model(**fun_input)
+
+    # Check that the output is the same
+    if isinstance(model_output_par_input, np.ndarray):
+        assert np.allclose(model_output_par_input, model_output_fun_input)
+    elif isinstance(model_output_par_input, cuqi.samples.Samples):
+        assert np.allclose(
+            model_output_par_input.samples, model_output_fun_input.samples
+        )
+
+
+@pytest.mark.parametrize("test_model, test_data", model_test_case_combinations)
+def test_gradient_of_multiple_input_model_is_correct(test_model, test_data):
+    """Test that the gradient method can handle multiple inputs and
+    return the correct output type and value"""
+
+    assert isinstance(test_model, cuqi.model.Model)
+    assert isinstance(test_data, TestCase)
+
+    # If case: gradient is not implemented for the model
+    if hasattr(test_model, "_do_test_gradient") and not test_model._do_test_gradient:
+        with pytest.raises(
+            NotImplementedError, match="Gradient is not implemented for this model."
+        ):
+            grad_output = test_model.gradient(
+                test_data.direction, **test_data.forward_input
+            )
+
+    # Elif case: gradient is implemented for the model and expected to work
+    elif not isinstance(
+        test_data.expected_grad_output_value,
+        (NotImplementedError, TypeError, ValueError),
+    ):
+        grad_output = test_model.gradient(
+            test_data.direction, **test_data.forward_input
+        )
+
+        # Assert output format is a dictionary with expected keys
+        assert list(grad_output.keys()) == test_model._non_default_args
+
+        # Check type and value of the output
+        for i, (k, v) in enumerate(grad_output.items()):
+            # Verify that the output is of the expected type
+            if v is not None:
+                assert isinstance(v, test_data.expected_grad_output_type)
+
+            # Verify that the output is of the expected value
+            if (
+                isinstance(test_model._gradient_func, tuple)
+                and test_model._gradient_func[i] is None
+            ):
+                assert v is None
+            else:
+                assert np.allclose(v, test_data.expected_grad_output_value[i])
+                assert np.allclose(v, test_data.FD_grad_output[i])
+
+    # Else case: gradient is implemented for the model but expected to raise an error
+    else:
+        with pytest.raises(
+            type(test_data.expected_grad_output_value),
+            match=str(test_data.expected_grad_output_value),
+        ):
+            grad_output = test_model.gradient(
+                test_data.direction, **test_data.forward_input
+            )
+
+
+# Create a sublist of model_test_case_combinations for cases where the expected
+# gradient output is not an error
+model_test_case_combinations_no_gradient_error = [
+    (test_model, test_data)
+    for test_model, test_data in model_test_case_combinations
+    if not isinstance(
+        test_data.expected_grad_output_value,
+        (NotImplementedError, TypeError, ValueError),
+    )
+]
+# Also, remove cases where the gradient is not implemented
+model_test_case_combinations_no_gradient_error = [
+    (test_model, test_data)
+    for test_model, test_data in model_test_case_combinations_no_gradient_error
+    if not hasattr(test_model, "_do_test_gradient") or test_model._do_test_gradient
+]
+
+
+@pytest.mark.parametrize(
+    "test_model, test_data", model_test_case_combinations_no_gradient_error
+)
+def test_gradient_of_multiple_input_model_accepts_stacked_input(test_model, test_data):
+    """Test that the gradient method can handle multiple inputs that are stacked and
+    return the correct output type and value"""
+
+    assert isinstance(test_model, cuqi.model.Model)
+    assert isinstance(test_data, TestCase)
+
+    grad_output_stacked_inputs = test_model.gradient(
+        test_data.direction, test_data.forward_input_stacked
+    )
+
+    # Assert output format is a dictionary with expected keys
+    assert list(grad_output_stacked_inputs.keys()) == test_model._non_default_args
+
+    # Assert type and value of the output are correct
+    for i, (k, v) in enumerate(grad_output_stacked_inputs.items()):
+        if v is not None:
+            assert isinstance(grad_output_stacked_inputs[k], np.ndarray)
+            assert np.allclose(
+                grad_output_stacked_inputs[k], test_data.expected_grad_output_value[i]
+            )
+
+    # Assert evaluating gradient on stacked input with wrong dimension raises error
+    with pytest.raises(
+        ValueError,
+        match=r"The number of positional arguments does not match the number of non-default arguments of the gradient. Additionally, the gradient input is specified by a single argument that cannot be split into multiple arguments matching the expected non_default_args",
+    ):
+        test_model.gradient(test_data.direction, test_data.forward_input_stacked[:-1])
+
+
+@pytest.mark.parametrize(
+    "test_model, test_data", model_test_case_combinations_no_gradient_error
+)
+def test_gradient_of_multiple_input_model_can_generate_stacked_output(
+    test_model, test_data
+):
+    """Test that the gradient method can generate stacked output of correct
+    value and type if the flag _gradient_output_stacked is enabled."""
+
+    # Generate stacked output for cases where the gradient specification is
+    # complete: all gradient functions are specified
+    if callable(test_model._gradient_func) or None not in test_model._gradient_func:
+        # Set the model flag _gradient_output_stacked to True to generate stacked output
+        test_model._gradient_output_stacked = True
+        stacked_grad_output = test_model.gradient(
+            test_data.direction, **test_data.forward_input
+        )
+        # Reset the flag to False
+        test_model._gradient_output_stacked = False
+
+        # Assert type and value of the output are correct
+        assert isinstance(stacked_grad_output, np.ndarray)
+        assert np.allclose(
+            stacked_grad_output, test_data.expected_staked_grad_output_value
+        )
+
+
+@pytest.mark.parametrize(
+    "test_model, test_data", model_test_case_combinations_no_gradient_error
+)
+def test_gradient_of_multiple_input_model_is_correct_with_funvals_input(
+    test_model, test_data
+):
+    """Test that the gradient method can handle multiple inputs some of which are
+    function values and return the correct output"""
+
+    # Par input
+    par_input = test_data.forward_input
+    grad_output_par_input = test_model.gradient(test_data.direction, **par_input)
+
+    # Fun input
+    fun_input = {
+        k: v.funvals if (isinstance(v, CUQIarray) or isinstance(v, Samples)) else v
+        for k, v in par_input.items()
+    }
+    grad_output_fun_input = test_model.gradient(test_data.direction, **fun_input)
+
+    # Check that the output is the same
+    for k, v in grad_output_par_input.items():
+        if v is not None:
+            assert np.allclose(v, grad_output_fun_input[k])
+        else:
+            assert grad_output_fun_input[k] is None
 
 
 @pytest.mark.parametrize("seed",[(0),(1),(2)])
@@ -19,7 +1333,7 @@ def test_LinearModel_getMatrix(seed):
     mat1 = model1.get_matrix() #Normal matrix
     mat2 = model2.get_matrix() #Sparse matrix (generated from functions)
 
-    assert np.allclose(mat1,mat2.A)
+    assert np.allclose(mat1,mat2.toarray())
 
 def test_initialize_model_dim():
     model1 = cuqi.model.Model(lambda x:x, range_geometry=4, domain_geometry=4)
@@ -71,7 +1385,7 @@ def test_forward(x, expected_type):
 
     model = cuqi.model.Model(forward=lambda x:A@x,
                             gradient=lambda direction,
-                            wrt: A.T@direction,
+                            x: A.T@direction,
                             domain_geometry=cuqi.geometry.Continuous1D(3),
                             range_geometry=cuqi.geometry.Continuous1D(2))
 
@@ -119,12 +1433,12 @@ def test_gradient(direction, expected_type):
 
     model = cuqi.model.Model(forward=lambda x:A@x,
                             gradient=lambda direction,
-                            wrt: A.T@direction,
+                            x: A.T@direction,
                             domain_geometry=cuqi.geometry.Continuous1D(3),
                             range_geometry=cuqi.geometry.Continuous1D(2))
 
     if isinstance(direction, cuqi.samples.Samples):
-        with pytest.raises(ValueError):
+        with pytest.raises(NotImplementedError, match="Gradient is not implemented for input of type Samples."):
            grad = model.gradient(direction, None)
     else:            
         grad = model.gradient(direction, None)
@@ -169,7 +1483,7 @@ def my_fun2par(p):
 
     model = cuqi.model.Model(forward=lambda x:x@A@x,
                             gradient=lambda direction,
-                            wrt: 2*wrt@A@direction,
+                            x: 2*x@A@direction,
                             domain_geometry=domain_geometry,
                             range_geometry=range_geometry)
 
@@ -178,33 +1492,33 @@ def my_fun2par(p):
         # the domain_geometry and hence can't compute fun2par 
         # (and wrt is passed as function).
         with pytest.raises(ValueError):
-            grad = model.gradient(direction, wrt, is_wrt_par=is_wrt_par)
+            grad = model.gradient(direction, wrt, is_var_par=is_wrt_par)
 
     if case_id ==2:
         # Raises an error because gradient is not implemented for
         # the domain_geometry and the domain geometry is not in
         # the list cuqi.geometry._get_identity_geometries()
         with pytest.raises(NotImplementedError):
-            grad = model.gradient(direction, wrt, is_wrt_par=is_wrt_par)
+            grad = model.gradient(direction, wrt, is_var_par=is_wrt_par)
 
     if case_id ==3:
         # Raises an error because domain_geometry does not have an
         # implementation of fun2par and wrt is passed as function.
         with pytest.raises(NotImplementedError):
-            grad = model.gradient(direction, wrt, is_wrt_par=is_wrt_par)
+            grad = model.gradient(direction, wrt, is_var_par=is_wrt_par)
 
     if case_id == 4:
         # Raises an error because the range_geometry is not in the
         # cuqi.geometry._get_identity_geometries() list
         with pytest.raises(NotImplementedError):
-            grad = model.gradient(direction, wrt, is_wrt_par=is_wrt_par)
+            grad = model.gradient(direction, wrt, is_var_par=is_wrt_par)
 
 
-@pytest.mark.parametrize("forward, gradient, direction, wrt, domain_geometry, domain_gradient, range_geometry, is_direction_par, is_wrt_par",
+@pytest.mark.parametrize("forward, gradient, direction, x, domain_geometry, domain_gradient, range_geometry, is_direction_par, is_wrt_par",
                          [
                              (
                                  lambda x: np.array([[1, 0, 0], [0, 3, .1]])@x,
-                                 lambda direction, wrt: np.array(
+                                 lambda direction, x: np.array(
                                      [[1, 0, 0], [0, 3, .1]]).T@direction,
                                  np.array([1, 12]),
                                  np.array([1, 12, 8]),
@@ -217,7 +1531,7 @@ def my_fun2par(p):
                              ),
                              (
                                  lambda x: np.array([[1, 2, 8], [1, 3, .1]])@x,
-                                 lambda direction, wrt: np.array(
+                                 lambda direction, x: np.array(
                                      [[1, 2, 8], [1, 3, .1]]).T@direction,
                                  np.array([1, 12]),
                                  np.array([1, 12, 8]),
@@ -311,7 +1625,7 @@ def my_fun2par(p):
                              ),
                          ]
                          )
-def test_gradient_computation(forward, gradient, direction, wrt, domain_geometry, domain_gradient, range_geometry, is_direction_par, is_wrt_par):
+def test_gradient_computation(forward, gradient, direction, x, domain_geometry, domain_gradient, range_geometry, is_direction_par, is_wrt_par):
     """Test applying chain rule when computing the gradient"""
     model = cuqi.model.Model(forward=forward,
                              gradient=gradient,
@@ -320,21 +1634,21 @@ def test_gradient_computation(forward, gradient, direction, wrt, domain_geometry
     model.domain_geometry.gradient = domain_gradient
 
     # Set CUQIarray geometries
-    if isinstance(wrt, CUQIarray):
-        wrt.geometry = domain_geometry
+    if isinstance(x, CUQIarray):
+        x.geometry = domain_geometry
     if isinstance(direction, CUQIarray):
         direction.geometry = range_geometry
     # Compute cuqi model gradient 
-    grad = model.gradient(direction, wrt, is_wrt_par=is_wrt_par)
+    grad = model.gradient(direction, x, is_var_par=is_wrt_par)
 
     # If wrt is function value, convert wrt to parameters to
     # be passed as an input for to cuqi.utilities.approx_derivative
-    if isinstance(wrt, CUQIarray):
-        wrt_ndarray = np.zeros(len(wrt.parameters))
-        wrt_ndarray[:] = wrt.parameters
-        wrt = wrt_ndarray
+    if isinstance(x, CUQIarray):
+        wrt_ndarray = np.zeros(len(x.parameters))
+        wrt_ndarray[:] = x.parameters
+        x = wrt_ndarray
     if not is_wrt_par:
-        wrt = domain_geometry.fun2par(wrt)
+        x = domain_geometry.fun2par(x)
 
     if isinstance(direction, CUQIarray):
         direction_ndarray = np.zeros(len(direction.parameters))
@@ -344,7 +1658,7 @@ def test_gradient_computation(forward, gradient, direction, wrt, domain_geometry
         direction = range_geometry.fun2par(direction)
     # Compute a finite difference approximation of cuqi model gradient 
     findiff_grad = cuqi.utilities.approx_derivative(
-        model.forward, wrt, direction)
+        model.forward, x, direction)
 
     # Compare the two gradients
     assert(np.allclose(grad, findiff_grad))
@@ -375,7 +1689,7 @@ def test_model_parameter_name_switch_errors():
 
     y = cuqi.distribution.Gaussian(np.zeros(model.domain_dim-1), 1)
 
-    with pytest.raises(ValueError, match=r"dimension does not match"):
+    with pytest.raises(ValueError, match=r"Attempting to match parameter name of Model with given distribution\(s\), but distribution\(s\) dimension\(s\) does not match model input dimension\(s\)."):
         model(y)    
 
 def test_model_allow_other_parameter_names():
@@ -406,45 +1720,73 @@ def forward_z(z, not_used=None):
     assert model_z(1) == 1
 
 def test_linear_model_allow_other_parameter_names():
-    """ Test that linear model automatically infers parameter names from the forward function and evaluation matches. """
+    """Test that linear model automatically infers parameter names from the forward and adjoint functions and evaluation matches."""
 
     forward_x = lambda x: x
     forward_y = lambda y: y
+
     def forward_z(z, not_used=None):
         return z
 
-    adjoint = lambda w: w
+    adjoint_w = lambda w: w
+    adjoint_v = lambda v: v
 
-    model_x = cuqi.model.LinearModel(forward_x, adjoint, 1, 1)
-    model_y = cuqi.model.LinearModel(forward_y, adjoint, 1, 1)
-    model_z = cuqi.model.LinearModel(forward_z, adjoint, 1, 1)
+    model_x_w = cuqi.model.LinearModel(forward_x, adjoint_w, 1, 1)
+    model_y_v = cuqi.model.LinearModel(forward_y, adjoint_v, 1, 1)
+    model_z_v = cuqi.model.LinearModel(forward_z, adjoint_v, 1, 1)
 
     A = np.array([[1]])
-    model_mat = cuqi.model.LinearModel(A) # Default parameter name is 'x'
+    model_mat = cuqi.model.LinearModel(A)  # Default parameter name is 'x'
 
+    # First: forward checks
     # Check that the model has switched its parameter name
-    assert model_x._non_default_args == ['x']
-    assert model_y._non_default_args == ['y']
-    assert model_z._non_default_args == ['z']
-    assert model_mat._non_default_args == ['x']
+    assert model_x_w._non_default_args == ["x"]
+    assert model_y_v._non_default_args == ["y"]
+    assert model_z_v._non_default_args == ["z"]
+    assert model_mat._non_default_args == ["x"]
 
     # Check that we can provide parameter names when evaluating the model
-    assert model_x(x=1) == 1
-    assert model_y(y=1) == 1
-    assert model_z(z=1) == 1
-    assert model_mat(x=np.ones(1)) == np.ones(1) # With matrix it has to be numpy array
+    assert model_x_w(x=1) == 1
+    assert model_y_v(y=1) == 1
+    assert model_z_v(z=1) == 1
+    assert model_mat(x=np.ones(1)) == np.ones(1)  # With matrix it has to be numpy array
+
+    # Check providing the wrong parameter name raises an error
+    with pytest.raises(
+        ValueError,
+        match=r"The model input is specified by a keywords arguments \['y'\] that does not match the non_default_args of the model \['x'\].",
+    ):
+        model_x_w(y=1)
 
     # And check that we can provide positional arguments
-    assert model_x(1) == 1
-    assert model_y(1) == 1
-    assert model_z(1) == 1
+    assert model_x_w(1) == 1
+    assert model_y_v(1) == 1
+    assert model_z_v(1) == 1
     assert model_mat(np.ones(1)) == np.ones(1)
 
     # Check that matrix multiplication works
-    assert model_x@1 == 1
-    assert model_y@1 == 1
-    assert model_z@1 == 1
-    assert model_mat@np.ones(1) == np.ones(1)
+    assert model_x_w @ 1 == 1
+    assert model_y_v @ 1 == 1
+    assert model_z_v @ 1 == 1
+    assert model_mat @ np.ones(1) == np.ones(1)
+
+    # Second: adjoint checks
+    # Check we can provide parameter names when evaluating the adjoint
+    assert model_x_w.adjoint(w=1) == 1
+    assert model_y_v.adjoint(v=1) == 1
+    assert model_z_v.adjoint(v=1) == 1
+
+    # Check providing the wrong parameter name raises an error
+    with pytest.raises(
+        ValueError,
+        match=r"The adjoint input is specified by a keywords arguments \['v'\] that does not match the non_default_args of the adjoint \['w'\].",
+    ):
+        model_x_w.adjoint(v=1)
+
+    # And check that we can provide positional arguments
+    assert model_x_w.adjoint(1) == 1
+    assert model_y_v.adjoint(1) == 1
+    assert model_z_v.adjoint(1) == 1
 
 def test_model_allows_jacobian_or_gradient():
     """ Test that either Jacobian or gradient (vector-jacobian product) can be specified and that it gives the same result. """
@@ -453,8 +1795,8 @@ def forward(x):
         return 10*x[1] - 10*x[0]**3 + 5*x[0]**2 + 6*x[0]
     def jacobian(x):
         return np.array([[-30*x[0]**2 + 10*x[0] + 6, 10]])
-    def gradient(dir, wrt):
-        return dir@jacobian(wrt)
+    def gradient(direction, x):
+        return direction@jacobian(x)
     
     # Check not both can be specified
     with pytest.raises(TypeError, match=r"Only one of gradient and jacobian"):
@@ -473,8 +1815,7 @@ def gradient(dir, wrt):
     assert np.allclose(model_grad.gradient(dir, wrt), model_jac.gradient(dir, wrt))
 
 # Parametrize over models
-@pytest.mark.parametrize("model", [cuqi.testproblem.Deconvolution1D().model,
-                                   cuqi.testproblem.Heat1D().model])
+@pytest.mark.parametrize("model", [cuqi.testproblem.Deconvolution1D().model])
 def test_AffineModel_Correct_result(model):
     """ Test creating a shifted linear model from a linear model """
 
@@ -490,6 +1831,47 @@ def test_AffineModel_Correct_result(model):
     # Check that the shifted linear model is correct
     assert np.allclose(A_affine(x), model(x) + b)
 
+@pytest.mark.parametrize("model", [cuqi.testproblem.Heat1D().model])
+def test_AffineModel_does_not_accept_general_model(model):
+    """Test that AffineModel only accepts LinearModel, besides accepting
+    a matrix or callable"""
+
+    # Random vectors
+    x = np.random.randn(model.domain_dim)
+    b = np.random.randn(model.range_dim)
+
+    with pytest.raises(
+        TypeError,
+        match=r"The linear operator should be a LinearModel object, a callable function or a matrix.",
+    ):
+        A_affine = cuqi.model.AffineModel(
+            model,
+            b,
+            range_geometry=model.range_geometry,
+            domain_geometry=model.domain_geometry,
+        )
+
+
+@pytest.mark.parametrize("fwd_func", [lambda: 1, lambda x, y: x])
+def test_AffineModel_raise_error_when_forward_or_adjoint_has_too_few_or_too_many_inputs(
+    fwd_func,
+):
+    """Test that AffineModel raises an error when the forward function or the adjoint function has too few or too many inputs"""
+    adj_func = lambda x: x
+
+    # Forward function with too few/many inputs
+    with pytest.raises(
+        ValueError, match=r"The linear operator should have exactly one input argument"
+    ):
+        cuqi.model.AffineModel(fwd_func, 0, adj_func)
+
+    # Adjoint function with too few/many inputs
+    with pytest.raises(
+        ValueError,
+        match=r"The adjoint linear operator should have exactly one input argument",
+    ):
+        cuqi.model.AffineModel(adj_func, 0, fwd_func)
+
 def test_AffineModel_update_shift():
 
     A = np.eye(2)
@@ -514,3 +1896,46 @@ def test_AffineModel_update_shift():
     assert np.all(model_copy(x) == np.array([3,0]))
 
 
+@pytest.mark.parametrize("domain_geometry, num_inputs, expected_domain_geometry", [
+    (4, 1, _DefaultGeometry1D(4)),
+    ((4,), 1, _DefaultGeometry1D(4)),
+    (Discrete(4), 1, Discrete(4)),
+    ((Discrete(4),), 1, Discrete(4)),
+    ((2, 2), 1, _DefaultGeometry2D((2, 2))),
+    (((2, 2),), 1, _DefaultGeometry2D((2, 2))),
+    (Image2D((2,2)), 1, Image2D((2,2))),
+    ((Image2D((2,2)),), 1, Image2D((2,2))),
+    ((Discrete(5), 3), 2, _ProductGeometry(Discrete(5), _DefaultGeometry1D(3))),
+    ((5, 3), 2, _ProductGeometry(_DefaultGeometry1D(5), _DefaultGeometry1D(3))),
+    ((5, Discrete(3),), 2, _ProductGeometry(_DefaultGeometry1D(5), Discrete(3)))
+])
+@pytest.mark.parametrize("range_geometry, expected_range_geometry", [
+    (1, _DefaultGeometry1D(1)),
+    ((1,), _DefaultGeometry1D(1)),
+    (Discrete(1), Discrete(1)),
+    ((Discrete(1),), Discrete(1))
+])
+def test_setting_domain_and_range_geometry(domain_geometry, range_geometry, num_inputs, expected_domain_geometry, expected_range_geometry):
+    """Test that the model can be initialized with different types of domain and range geometries"""
+    # Set up the forward map for 1 input
+    forward = lambda x: np.sum(x)
+    # Set up the forward map for 2 inputs
+    forward_2_inputs = lambda x, y: np.sum(x) + np.sum(y)
+
+    # Create the model with the specified domain and range geometries
+    if num_inputs == 1:
+        model = cuqi.model.Model(
+            forward,
+            domain_geometry=domain_geometry,
+            range_geometry=range_geometry,
+        )
+    else:
+        model = cuqi.model.Model(
+            forward_2_inputs,
+            domain_geometry=domain_geometry,
+            range_geometry=range_geometry,
+        )
+
+    # Check the model domain and range geometries are of the expected type
+    assert model.domain_geometry == expected_domain_geometry
+    assert model.range_geometry == expected_range_geometry
\ No newline at end of file
diff --git a/tests/test_solver.py b/tests/test_solver.py
index 62e35cfc2..8987fb760 100644
--- a/tests/test_solver.py
+++ b/tests/test_solver.py
@@ -100,19 +100,27 @@ def test_ScipyLinearLSQ_with_LinearOperator():
     sol, _ = ScipyLinearLSQ(A_op, b).solve()
     assert np.allclose(sol, ref_sol, rtol=1e-10)
 
-def test_ScipyLinearLSQ_against_FISTA():
+def test_ScipyLinearLSQ_against_ScipyMinimizer_and_against_FISTA():
     A = np.array([[73,71,52],[87,74,46],[72,2,7],[80,89,71]])
     b = np.array([49,67,68,20])
+
+    # solve with ScipyMinimizer
+    def fun(x):
+        return 0.5*np.linalg.norm(A@x-b)**2
+    def jac(x):
+        return A.T@(A@x-b)
+    sol_min, _ = ScipyMinimizer(fun, np.zeros(3), gradfunc=jac, tol=1e-10, bounds=[(0,np.inf),(0,np.inf),(0,np.inf)]).solve()
+
     # solve with ScipyLinearLSQ
-    lb = np.zeros(3)
-    ub = lb + np.inf
-    sol_lsq, _ = ScipyLinearLSQ(A, b, (lb,ub)).solve()
+    sol_lsq, _ = ScipyLinearLSQ(A, b, ([0,0,0],[np.inf,np.inf,np.inf]), tol=1e-10).solve()
+
     # solve with FISTA
     rng = np.random.default_rng(seed = 1219)
     x0 = rng.standard_normal(3)
     sol_fista, _ = FISTA(A, b, lambda x, _: ProjectNonnegative(x), x0, stepsize=1e-7, maxit=100000, abstol=1e-16, adaptive=True).solve()
 
     assert np.allclose(sol_lsq, sol_fista, rtol=1e-8)
+    assert np.allclose(sol_min, sol_lsq, rtol=1e-8)
 
 def test_LM():
     # compare to MATLAB's original code solution
diff --git a/tests/test_utilities.py b/tests/test_utilities.py
index 0959f550d..9e74ff5fb 100644
--- a/tests/test_utilities.py
+++ b/tests/test_utilities.py
@@ -1,7 +1,7 @@
 import pytest
 from scipy.linalg import cholesky
 from scipy.sparse import diags
-from cuqi.utilities import sparse_cholesky, plot_1D_density, plot_2D_density, count_nonzero, count_constant_components_1D, count_constant_components_2D
+from cuqi.utilities import sparse_cholesky, plot_1D_density, plot_2D_density, count_nonzero, count_within_bounds, count_constant_components_1D, count_constant_components_2D, piecewise_linear_1D_DoF
 from cuqi.model import LinearModel
 from cuqi.distribution import Gaussian, Uniform, JointDistribution
 import numpy as np
@@ -72,11 +72,19 @@ def count_nonzero_threshold():
     eps_min = 1e-1
     eps_plus = 1e1
     test_array = np.array([0, -eps_min*threshold, -threshold, -eps_plus*threshold, eps_min*threshold, threshold, eps_plus*threshold])
-    print(test_array)
 
     nonzeros = count_nonzero(test_array, threshold)
     assert(nonzeros == 4)
 
+
+def count_within_bounds_exceptions():
+    test_array = np.array([-1.0, -0.5, 0.0, 0.5, 1.0])
+    lower = -np.ones_like(test_array)
+    upper = np.ones_like(test_array)
+    
+    within_bounds = count_within_bounds(test_array, lower, upper, exception=0.0)
+    assert(within_bounds == 2)
+
 def count_constant_components_1d_threshold():
     threshold = 1e-4
     test_array = np.array([0, 0, 1.0 - 0.5*threshold, 1.0, 1.0+0.5*threshold, 1.5, 2, 2+0.5*threshold, 2+threshold, 0, 0])
@@ -87,6 +95,9 @@ def count_constant_components_1d_threshold():
     components = count_constant_components_1D(test_array, threshold, lower = 0, upper = 2)
     assert(components == 2)
 
+    components = count_constant_components_1D(test_array, threshold, exception = 0)
+    assert(components == 3)
+
 def count_constant_components_2d_threshold():
     threshold = 1e-4
     test_array = np.array([[0, 0                , 1 - 0.25*threshold],
@@ -97,4 +108,18 @@ def count_constant_components_2d_threshold():
     assert(components == 3)
 
     components = count_constant_components_2D(test_array, threshold, lower = 0, upper = 2)
-    assert(components == 1)
\ No newline at end of file
+    assert(components == 1)
+
+    
+def piecewise_linear_1D_DoF_exceptions():
+    threshold = 1e-4
+    test_array = np.array([0, 0, 0, 0.5, 1.0, 1.5, 2.0, 2.0, 2.0, 3.0, 4.0])
+
+    components = piecewise_linear_1D_DoF(test_array, threshold)
+    assert(components == 5)
+
+    components = piecewise_linear_1D_DoF(test_array, threshold, exception_zero = True)
+    assert(components == 3)
+
+    components = piecewise_linear_1D_DoF(test_array, threshold, exception_flat = True)
+    assert(components == 3)
\ No newline at end of file
diff --git a/tests/zexperimental/test_mcmc.py b/tests/zexperimental/test_mcmc.py
index eb60fd557..8a9f0d563 100644
--- a/tests/zexperimental/test_mcmc.py
+++ b/tests/zexperimental/test_mcmc.py
@@ -1510,3 +1510,126 @@ def test_RegularizedLinearRTO_ScipyLinearLSQ_option_invalid():
 
     with pytest.raises(ValueError, match="ScipyLinearLSQ"):
         sampler = cuqi.experimental.mcmc.RegularizedLinearRTO(posterior, solver = "ScipyLinearLSQ")
+
+def test_RegularizedLinearRTO_ScipyLinearLSQ_against_ScipyMinimizer_and_against_FISTA():
+    # Define LinearModel and data
+    A, y_obs, _ = cuqi.testproblem.Deconvolution1D().get_components()
+
+    # Define Bayesian Problem
+    x = cuqi.implicitprior.NonnegativeGMRF(np.zeros(A.domain_dim), 100)
+    y = cuqi.distribution.Gaussian(A@x, 0.01**2)
+    posterior = cuqi.distribution.JointDistribution(x, y)(y=y_obs)
+
+    # Set up RegularizedLinearRTO with three solvers
+    sampler1 = cuqi.experimental.mcmc.RegularizedLinearRTO(posterior, solver="ScipyMinimizer", maxit=1000, tol=1e-8)
+    sampler2 = cuqi.experimental.mcmc.RegularizedLinearRTO(posterior, solver="ScipyLinearLSQ", maxit=1000, tol=1e-8)
+    sampler3 = cuqi.experimental.mcmc.RegularizedLinearRTO(posterior, solver="FISTA", maxit=1000, tol=1e-8)
+
+    # Sample with fixed seed
+    np.random.seed(0)
+    samples1 = sampler1.sample(5).get_samples()
+    np.random.seed(0)
+    samples2 = sampler2.sample(5).get_samples()
+    np.random.seed(0)
+    samples3 = sampler3.sample(5).get_samples()
+
+    assert np.allclose(samples1.samples.mean(), samples2.samples.mean(), rtol=1e-5)
+    assert np.allclose(samples1.samples.mean(), samples3.samples.mean(), rtol=1e-5)
+
+# ============ Start testing sampler callback ============
+# Samplers that should be tested for callback
+callback_testing_sampler_classes = [
+    cls
+    for _, cls in inspect.getmembers(cuqi.experimental.mcmc, inspect.isclass)
+    if cls not in [cuqi.experimental.mcmc.Sampler, cuqi.experimental.mcmc.ProposalBasedSampler]
+]
+
+# Instances of samplers that should be tested for callback
+callback_testing_sampler_instances = [
+    cuqi.experimental.mcmc.MH(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.PCN(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.CWMH(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.ULA(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.MALA(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.NUTS(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.LinearRTO(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.RegularizedLinearRTO(
+        target=create_regularized_target(dim=16)
+    ),
+    cuqi.experimental.mcmc.UGLA(target=create_lmrf_prior_target(dim=16)),
+    cuqi.experimental.mcmc.Direct(target=cuqi.distribution.Gaussian(np.zeros(10), 1)),
+    cuqi.experimental.mcmc.Conjugate(target=create_conjugate_target("Gaussian-Gamma")),
+    cuqi.experimental.mcmc.ConjugateApprox(
+        target=create_conjugate_target("LMRF-Gamma")
+    ),
+    cuqi.experimental.mcmc.MYULA(target=create_myula_target(dim=16)),
+    cuqi.experimental.mcmc.PnPULA(target=create_myula_target(dim=16)),
+    cuqi.experimental.mcmc.HybridGibbs(
+        target=HybridGibbs_target_1(),
+        sampling_strategy={
+            "x" : cuqi.experimental.mcmc.NUTS(max_depth=7),
+            "s" : cuqi.experimental.mcmc.Conjugate()
+        },
+    ),
+]
+
+@pytest.mark.parametrize("sampler", callback_testing_sampler_instances)
+def test_passing_callback_to_sampler(sampler):
+    """ Test that the callback function is called by the sampler and
+    the sampler passes the correct arguments to the callback function. """
+
+    Ns = 10
+    indices = [] # list to store the indices of the samples
+
+    # Define the callback function
+    def callback_function(callback_sampler, sample_index, num_samples):
+        assert isinstance(callback_sampler, sampler.__class__)
+        assert num_samples == Ns
+        indices.append(sample_index)
+
+    # Create a sampler with callback function
+    if not isinstance(sampler, cuqi.experimental.mcmc.HybridGibbs):
+        my_sampler = sampler.__class__(sampler.target, callback=callback_function)
+    else:
+        my_sampler = cuqi.experimental.mcmc.HybridGibbs(
+            sampler.target,
+            sampling_strategy={
+                "x": cuqi.experimental.mcmc.NUTS(max_depth=7),
+                "s": cuqi.experimental.mcmc.Conjugate(),
+            },
+            callback=callback_function,
+        )
+
+    # Sample Ns warmup samples
+    my_sampler.warmup(Ns)
+    assert np.allclose(indices, np.arange(Ns))
+
+    # reset the indices
+    indices = []
+
+    # Sample Ns samples
+    my_sampler.sample(Ns)
+    assert np.allclose(indices, np.arange(Ns))
+
+def test_all_samplers_that_should_be_tested_for_callback_are_in_the_tested_list():
+    """ Test that all samplers that should be tested for callback are in the callback_testing_sampler_instances. """
+    # The classes of the tested samplers:
+    tested_classes = [sampler.__class__ for sampler in callback_testing_sampler_instances]
+    for cls in callback_testing_sampler_classes:
+        assert cls in tested_classes, f"Sampler {cls} is not tested for callback."
+
+# ============= End testing sampler callback =============
\ No newline at end of file
diff --git a/tests/zexperimental/test_randomvariable.py b/tests/zexperimental/test_randomvariable.py
index cbd1f5831..a0ff54bb0 100644
--- a/tests/zexperimental/test_randomvariable.py
+++ b/tests/zexperimental/test_randomvariable.py
@@ -62,8 +62,8 @@ def test_algrabraic_operations_on_distribution_should_create_randomvariable():
     lambda x: cuqi.model.LinearModel(np.ones((2,2)))@x,
 ])
 def test_algebra_on_randomvariables_can_be_computed_and_is_correct(operations):
-    X = RandomVariable(cuqi.distribution.Gaussian(np.zeros(2), 1))
-    rv = operations(X)
+    x = RandomVariable(cuqi.distribution.Gaussian(np.zeros(2), 1))
+    rv = operations(x)
     val = np.random.randn(2)
     # Compare random variable recorded operations vs actual operations
     assert np.allclose(rv(val), operations(val))