Repository: pymc-devs/pymc

Issue:
ENH: Add a way to specify `mode` for `check_start_vals`
### Before

_No response_

### After

```python
with pm.Model() as m:
    ...
m.check_start_vals(..., mode="JAX")
```


### Context for the issue:

Sometimes we deploy `pymc` in environments were we actually only care about GPU sampling via JAX, so there shouldn't be a need for a C compiler. This is all fine and good because we can `sample_prior_predictive`, `sample` and `sample_posterior_predictive` using JAX or any other mode that doesn't need C. The problem comes when `sample` calls `model.check_start_vals`, this will use the default mode and in the absence of a compiler, the python code might take very long to run. It would be great to have a way to set the compilation mode for `check_start_vals` as well.

Discussion:
ricardoV94: I changed the label to maintenance to de-emphasize in the release notes. Feel free to revert if you disagree 
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pymc-devs/pymc/pull/7482?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 92.16%. Comparing base [(`c92a9a9`)](https://app.codecov.io/gh/pymc-devs/pymc/commit/c92a9a9463f501d5e7977fd80719df6139425322?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) to head [(`5db847d`)](https://app.codecov.io/gh/pymc-devs/pymc/commit/5db847d043ab61468617e782b0b78b0e7c71c618?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs).

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://app.codecov.io/gh/pymc-devs/pymc/pull/7482/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://app.codecov.io/gh/pymc-devs/pymc/pull/7482?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@           Coverage Diff           @@
##             main    #7482   +/-   ##
=======================================
  Coverage   92.16%   92.16%           
=======================================
  Files         103      103           
  Lines       17214    17214           
=======================================
  Hits        15866    15866           
  Misses       1348     1348           
```

| [Files](https://app.codecov.io/gh/pymc-devs/pymc/pull/7482?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/model/core.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7482?src=pr&el=tree&filepath=pymc%2Fmodel%2Fcore.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9tb2RlbC9jb3JlLnB5) | `91.77% <100.00%> (ø)` | |

</details>

Patch:
diff --git a/pymc/model/core.py b/pymc/model/core.py
index 3a2741766..71dcfee48 100644
--- a/pymc/model/core.py
+++ b/pymc/model/core.py
@@ -1652,6 +1652,8 @@ class Model(WithMemoization, metaclass=ContextMeta):
         point_fn : bool
             Whether to wrap the compiled function in a PointFunc, which takes a Point
             dictionary with model variable names and values as input.
+        Other keyword arguments :
+            Any other keyword argument is sent to :py:func:`pymc.pytensorf.compile_pymc`.
 
         Returns
         -------
@@ -1747,7 +1749,7 @@ class Model(WithMemoization, metaclass=ContextMeta):
         )
         return {name: tuple(shape) for name, shape in zip(names, f())}
 
-    def check_start_vals(self, start):
+    def check_start_vals(self, start, **kwargs):
         r"""Check that the starting values for MCMC do not cause the relevant log probability
         to evaluate to something invalid (e.g. Inf or NaN)
 
@@ -1758,6 +1760,8 @@ class Model(WithMemoization, metaclass=ContextMeta):
             Defaults to ``trace.point(-1))`` if there is a trace provided and
             ``model.initial_point`` if not (defaults to empty dict). Initialization
             methods for NUTS (see ``init`` keyword) can overwrite the default.
+        Other keyword arguments :
+            Any other keyword argument is sent to :py:meth:`~pymc.model.core.Model.point_logps`.
 
         Raises
         ------
@@ -1787,7 +1791,7 @@ class Model(WithMemoization, metaclass=ContextMeta):
                     f"Valid keys are: {valid_keys}, but {extra_keys} was supplied"
                 )
 
-            initial_eval = self.point_logps(point=elem)
+            initial_eval = self.point_logps(point=elem, **kwargs)
 
             if not all(np.isfinite(v) for v in initial_eval.values()):
                 raise SamplingError(
@@ -1797,7 +1801,7 @@ class Model(WithMemoization, metaclass=ContextMeta):
                     "You can call `model.debug()` for more details."
                 )
 
-    def point_logps(self, point=None, round_vals=2):
+    def point_logps(self, point=None, round_vals=2, **kwargs):
         """Computes the log probability of `point` for all random variables in the model.
 
         Parameters
@@ -1807,6 +1811,8 @@ class Model(WithMemoization, metaclass=ContextMeta):
             is used.
         round_vals : int, default 2
             Number of decimals to round log-probabilities.
+        Other keyword arguments :
+            Any other keyword argument are sent provided to :py:meth:`~pymc.model.core.Model.compile_fn`
 
         Returns
         -------
@@ -1822,7 +1828,7 @@ class Model(WithMemoization, metaclass=ContextMeta):
             factor.name: np.round(np.asarray(factor_logp), round_vals)
             for factor, factor_logp in zip(
                 factors,
-                self.compile_fn(factor_logps_fn)(point),
+                self.compile_fn(factor_logps_fn, **kwargs)(point),
             )
         }
 


Test patch:
diff --git a/tests/model/test_core.py b/tests/model/test_core.py
index 669704bb5..2504f5c79 100644
--- a/tests/model/test_core.py
+++ b/tests/model/test_core.py
@@ -756,6 +756,20 @@ class TestCheckStartVals:
         with pytest.raises(KeyError):
             model.check_start_vals(start)
 
+    @pytest.mark.parametrize("mode", [None, "JAX", "NUMBA"])
+    def test_mode(self, mode):
+        with pm.Model() as model:
+            a = pm.Uniform("a", lower=0.0, upper=1.0)
+            b = pm.Uniform("b", lower=2.0, upper=3.0)
+        start = {
+            "a_interval__": model.rvs_to_transforms[a].forward(0.3, *a.owner.inputs).eval(),
+            "b_interval__": model.rvs_to_transforms[b].forward(2.1, *b.owner.inputs).eval(),
+        }
+        with patch("pymc.model.core.compile_pymc") as patched_compile_pymc:
+            model.check_start_vals(start, mode=mode)
+        patched_compile_pymc.assert_called_once()
+        assert patched_compile_pymc.call_args.kwargs["mode"] == mode
+
 
 def test_set_initval():
     # Make sure the dependencies between variables are maintained when