ValueGradFunction inner function now accepts a raveled input

Author: ricardoV94

pymc/model/core.py

@@ -61,6 +61,7 @@
     gradient,
     hessian,
     inputvars,
+    join_nonshared_inputs,
     rewrite_pregrad,
 )
 from pymc.util import (
@@ -172,6 +173,9 @@ def __init__(
         dtype=None,
         casting="no",
         compute_grads=True,
+        model=None,
+        initial_point=None,
+        ravel_inputs: bool | None = None,
         **kwargs,
     ):
         if extra_vars_and_values is None:
@@ -219,9 +223,7 @@ def __init__(
         givens = []
         self._extra_vars_shared = {}
         for var, value in extra_vars_and_values.items():
-            shared = pytensor.shared(
-                value, var.name + "_shared__", shape=[1 if s == 1 else None for s in value.shape]
-            )
+            shared = pytensor.shared(value, var.name + "_shared__", shape=value.shape)
             self._extra_vars_shared[var.name] = shared
             givens.append((var, shared))
 
@@ -231,13 +233,28 @@ def __init__(
             grads = pytensor.grad(cost, grad_vars, disconnected_inputs="ignore")
             for grad_wrt, var in zip(grads, grad_vars):
                 grad_wrt.name = f"{var.name}_grad"
-            outputs = [cost, *grads]
+            grads = pt.join(0, *[pt.atleast_1d(grad.ravel()) for grad in grads])
+            outputs = [cost, grads]
         else:
             outputs = [cost]
 
-        inputs = grad_vars
+        if ravel_inputs:
+            if initial_point is None:
+                initial_point = modelcontext(model).initial_point()
+            outputs, raveled_grad_vars = join_nonshared_inputs(
+                point=initial_point, inputs=grad_vars, outputs=outputs, make_inputs_shared=False
+            )
+            inputs = [raveled_grad_vars]
+        else:
+            if ravel_inputs is None:
+                warnings.warn(
+                    "ValueGradFunction will become a function of raveled inputs.\n"
+                    "Specify `ravel_inputs` to suppress this warning. Note that setting `ravel_inputs=False` will be forbidden in a future release."
+                )
+            inputs = grad_vars
 
         self._pytensor_function = compile_pymc(inputs, outputs, givens=givens, **kwargs)
+        self._raveled_inputs = ravel_inputs
 
     def set_weights(self, values):
         if values.shape != (self._n_costs - 1,):
@@ -247,38 +264,29 @@ def set_weights(self, values):
     def set_extra_values(self, extra_vars):
         self._extra_are_set = True
         for var in self._extra_vars:
-            self._extra_vars_shared[var.name].set_value(extra_vars[var.name])
+            self._extra_vars_shared[var.name].set_value(extra_vars[var.name], borrow=True)
 
     def get_extra_values(self):
         if not self._extra_are_set:
             raise ValueError("Extra values are not set.")
 
         return {var.name: self._extra_vars_shared[var.name].get_value() for var in self._extra_vars}
 
-    def __call__(self, grad_vars, grad_out=None, extra_vars=None):
+    def __call__(self, grad_vars, *, extra_vars=None):
         if extra_vars is not None:
             self.set_extra_values(extra_vars)
-
-        if not self._extra_are_set:
+        elif not self._extra_are_set:
             raise ValueError("Extra values are not set.")
 
         if isinstance(grad_vars, RaveledVars):
-            grad_vars = list(DictToArrayBijection.rmap(grad_vars).values())
-
-        cost, *grads = self._pytensor_function(*grad_vars)
-
-        if grads:
-            grads_raveled = DictToArrayBijection.map(
-                {v.name: gv for v, gv in zip(self._grad_vars, grads)}
-            )
-
-            if grad_out is None:
-                return cost, grads_raveled.data
+            if self._raveled_inputs:
+                grad_vars = (grad_vars.data,)
             else:
-                np.copyto(grad_out, grads_raveled.data)
-                return cost
-        else:
-            return cost
+                grad_vars = DictToArrayBijection.rmap(grad_vars).values()
+        elif self._raveled_inputs and not isinstance(grad_vars, Sequence):
+            grad_vars = (grad_vars,)
+
+        return self._pytensor_function(*grad_vars)
 
     @property
     def profile(self):
@@ -521,7 +529,14 @@ def root(self):
     def isroot(self):
         return self.parent is None
 
-    def logp_dlogp_function(self, grad_vars=None, tempered=False, **kwargs):
+    def logp_dlogp_function(
+        self,
+        grad_vars=None,
+        tempered=False,
+        initial_point=None,
+        ravel_inputs: bool | None = None,
+        **kwargs,
+    ):
         """Compile a PyTensor function that computes logp and gradient.
 
         Parameters
@@ -547,13 +562,22 @@ def logp_dlogp_function(self, grad_vars=None, tempered=False, **kwargs):
             costs = [self.logp()]
 
         input_vars = {i for i in graph_inputs(costs) if not isinstance(i, Constant)}
-        ip = self.initial_point(0)
+        if initial_point is None:
+            initial_point = self.initial_point(0)
         extra_vars_and_values = {
-            var: ip[var.name]
+            var: initial_point[var.name]
             for var in self.value_vars
             if var in input_vars and var not in grad_vars
         }
-        return ValueGradFunction(costs, grad_vars, extra_vars_and_values, **kwargs)
+        return ValueGradFunction(
+            costs,
+            grad_vars,
+            extra_vars_and_values,
+            model=self,
+            initial_point=initial_point,
+            ravel_inputs=ravel_inputs,
+            **kwargs,
+        )
 
     def compile_logp(
         self,

pymc/sampling/mcmc.py

@@ -1441,6 +1441,8 @@ def init_nuts(
         pm.callbacks.CheckParametersConvergence(tolerance=1e-2, diff="relative"),
     ]
 
+    logp_dlogp_func = model.logp_dlogp_function(ravel_inputs=True)
+    logp_dlogp_func.trust_input = True
     initial_points = _init_jitter(
         model,
         initvals,

pymc/step_methods/arraystep.py

@@ -185,17 +185,17 @@ def __init__(
         model = modelcontext(model)
 
         if logp_dlogp_func is None:
-            func = model.logp_dlogp_function(vars, dtype=dtype, **pytensor_kwargs)
-        else:
-            func = logp_dlogp_func
-
-        self._logp_dlogp_func = func
+            logp_dlogp_func = model.logp_dlogp_function(
+                vars,
+                dtype=dtype,
+                ravel_inputs=True,
+                **pytensor_kwargs,
+            )
+            logp_dlogp_func.trust_input = True
 
-        super().__init__(vars, func._extra_vars_shared, blocked, rng=rng)
+        self._logp_dlogp_func = logp_dlogp_func
 
-    def step(self, point) -> tuple[PointType, StatsType]:
-        self._logp_dlogp_func._extra_are_set = True
-        return super().step(point)
+        super().__init__(vars, logp_dlogp_func._extra_vars_shared, blocked, rng=rng)
 
 
 def metrop_select(

pymc/step_methods/hmc/base_hmc.py

@@ -194,8 +194,6 @@ def astep(self, q0: RaveledVars) -> tuple[RaveledVars, StatsType]:
         process_start = time.process_time()
 
         p0 = self.potential.random()
-        p0 = RaveledVars(p0, q0.point_map_info)
-
         start = self.integrator.compute_state(q0, p0)
 
         warning: SamplerWarning | None = None
@@ -226,13 +224,13 @@ def astep(self, q0: RaveledVars) -> tuple[RaveledVars, StatsType]:
         if self._step_rand is not None:
             step_size = self._step_rand(step_size, rng=self.rng)
 
-        hmc_step = self._hamiltonian_step(start, p0.data, step_size)
+        hmc_step = self._hamiltonian_step(start, p0, step_size)
 
         perf_end = time.perf_counter()
         process_end = time.process_time()
 
         self.step_adapt.update(hmc_step.accept_stat, adapt_step)
-        self.potential.update(hmc_step.end.q, hmc_step.end.q_grad, self.tune)
+        self.potential.update(hmc_step.end.q.data, hmc_step.end.q_grad, self.tune)
         if hmc_step.divergence_info:
             info = hmc_step.divergence_info
             point = None

pymc/step_methods/hmc/integration.py

@@ -18,13 +18,13 @@
 
 from scipy import linalg
 
-from pymc.blocking import RaveledVars
+from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.step_methods.hmc.quadpotential import QuadPotential
 
 
 class State(NamedTuple):
     q: RaveledVars
-    p: RaveledVars
+    p: np.ndarray
     v: np.ndarray
     q_grad: np.ndarray
     energy: float
@@ -40,23 +40,35 @@ class CpuLeapfrogIntegrator:
     def __init__(self, potential: QuadPotential, logp_dlogp_func):
         """Leapfrog integrator using CPU."""
         self._potential = potential
-        self._logp_dlogp_func = logp_dlogp_func
-        self._dtype = self._logp_dlogp_func.dtype
+        # Sidestep logp_dlogp_function.__call__
+        pytensor_function = logp_dlogp_func._pytensor_function
+        # Create some wrappers for backwards compatibility during transition
+        # When raveled_inputs=False is forbidden, func = pytensor_function
+        if logp_dlogp_func._raveled_inputs:
+
+            def func(q, _):
+                return pytensor_function(q)
+
+        else:
+
+            def func(q, point_map_info):
+                unraveled_q = DictToArrayBijection.rmap(RaveledVars(q, point_map_info)).values()
+                return pytensor_function(*unraveled_q)
+
+        self._logp_dlogp_func = func
+        self._dtype = logp_dlogp_func.dtype
         if self._potential.dtype != self._dtype:
             raise ValueError(
                 f"dtypes of potential ({self._potential.dtype}) and logp function ({self._dtype})"
                 "don't match."
             )
 
-    def compute_state(self, q: RaveledVars, p: RaveledVars):
+    def compute_state(self, q: RaveledVars, p: np.ndarray):
         """Compute Hamiltonian functions using a position and momentum."""
-        if q.data.dtype != self._dtype or p.data.dtype != self._dtype:
-            raise ValueError(f"Invalid dtype. Must be {self._dtype}")
-
-        logp, dlogp = self._logp_dlogp_func(q)
+        logp, dlogp = self._logp_dlogp_func(q.data, q.point_map_info)
 
-        v = self._potential.velocity(p.data, out=None)
-        kinetic = self._potential.energy(p.data, velocity=v)
+        v = self._potential.velocity(p, out=None)
+        kinetic = self._potential.energy(p, velocity=v)
         energy = kinetic - logp
         return State(q, p, v, dlogp, energy, logp, 0)
 
@@ -96,10 +108,10 @@ def _step(self, epsilon, state):
         axpy = linalg.blas.get_blas_funcs("axpy", dtype=self._dtype)
         pot = self._potential
 
-        q_new = state.q.data.copy()
-        p_new = state.p.data.copy()
+        q = state.q
+        q_new = q.data.copy()
+        p_new = state.p.copy()
         v_new = np.empty_like(q_new)
-        q_new_grad = np.empty_like(q_new)
 
         dt = 0.5 * epsilon
 
@@ -112,19 +124,16 @@ def _step(self, epsilon, state):
         # q_new = q + epsilon * v_new
         axpy(v_new, q_new, a=epsilon)
 
-        p_new = RaveledVars(p_new, state.p.point_map_info)
-        q_new = RaveledVars(q_new, state.q.point_map_info)
-
-        logp = self._logp_dlogp_func(q_new, grad_out=q_new_grad)
+        logp, q_new_grad = self._logp_dlogp_func(q_new, q.point_map_info)
 
         # p_new = p_new + dt * q_new_grad
-        axpy(q_new_grad, p_new.data, a=dt)
+        axpy(q_new_grad, p_new, a=dt)
 
-        kinetic = pot.velocity_energy(p_new.data, v_new)
+        kinetic = pot.velocity_energy(p_new, v_new)
         energy = kinetic - logp
 
         return State(
-            q_new,
+            RaveledVars(q_new, state.q.point_map_info),
             p_new,
             v_new,
             q_new_grad,

pymc/step_methods/hmc/nuts.py

@@ -279,7 +279,7 @@ def __init__(
         self.log_accept_sum = -np.inf
         self.mean_tree_accept = 0.0
         self.n_proposals = 0
-        self.p_sum = start.p.data.copy()
+        self.p_sum = start.p.copy()
         self.max_energy_change = 0.0
 
     def extend(self, direction):
@@ -330,9 +330,9 @@ def extend(self, direction):
             left, right = self.left, self.right
             p_sum = self.p_sum
             turning = (p_sum.dot(left.v) <= 0) or (p_sum.dot(right.v) <= 0)
-            p_sum1 = leftmost_p_sum + rightmost_begin.p.data
+            p_sum1 = leftmost_p_sum + rightmost_begin.p
             turning1 = (p_sum1.dot(leftmost_begin.v) <= 0) or (p_sum1.dot(rightmost_begin.v) <= 0)
-            p_sum2 = leftmost_end.p.data + rightmost_p_sum
+            p_sum2 = leftmost_end.p + rightmost_p_sum
             turning2 = (p_sum2.dot(leftmost_end.v) <= 0) or (p_sum2.dot(rightmost_end.v) <= 0)
             turning = turning | turning1 | turning2
 
@@ -372,7 +372,7 @@ def _single_step(self, left: State, epsilon: float):
                     right.model_logp,
                     right.index_in_trajectory,
                 )
-                tree = Subtree(right, right, right.p.data, proposal, log_size)
+                tree = Subtree(right, right, right.p, proposal, log_size)
                 return tree, None, False
             else:
                 error_msg = f"Energy change in leapfrog step is too large: {energy_change}."
@@ -400,9 +400,9 @@ def _build_subtree(self, left, depth, epsilon):
             turning = (p_sum.dot(left.v) <= 0) or (p_sum.dot(right.v) <= 0)
             # Additional U turn check only when depth > 1 to avoid redundant work.
             if depth - 1 > 0:
-                p_sum1 = tree1.p_sum + tree2.left.p.data
+                p_sum1 = tree1.p_sum + tree2.left.p
                 turning1 = (p_sum1.dot(tree1.left.v) <= 0) or (p_sum1.dot(tree2.left.v) <= 0)
-                p_sum2 = tree1.right.p.data + tree2.p_sum
+                p_sum2 = tree1.right.p + tree2.p_sum
                 turning2 = (p_sum2.dot(tree1.right.v) <= 0) or (p_sum2.dot(tree2.right.v) <= 0)
                 turning = turning | turning1 | turning2
 

tests/distributions/test_multivariate.py

@@ -2395,7 +2395,7 @@ def test_mvnormal_no_cholesky_in_model_logp():
     d2logp = m.compile_d2logp()
     assert not contains_cholesky_op(d2logp.f.maker.fgraph)
 
-    logp_dlogp = m.logp_dlogp_function()
+    logp_dlogp = m.logp_dlogp_function(ravel_inputs=True)
     assert not contains_cholesky_op(logp_dlogp._pytensor_function.maker.fgraph)