Avoid input copy in Ndarray fn

Also initialize empty trace and set `trust_input=True`

Author: ricardoV94

pymc/backends/base.py

@@ -30,9 +30,11 @@
 )
 
 import numpy as np
+import pytensor
 
 from pymc.backends.report import SamplerReport
 from pymc.model import modelcontext
+from pymc.pytensorf import compile_pymc
 from pymc.util import get_var_name
 
 logger = logging.getLogger(__name__)
@@ -168,7 +170,13 @@ def __init__(
             raise Exception(f"Can't trace unnamed variables: {unnamed_vars}")
 
         if fn is None:
-            fn = model.compile_fn(vars, inputs=model.value_vars, on_unused_input="ignore")
+            # borrow=True avoids deepcopy when inputs=output which is the case for untransformed value variables
+            fn = compile_pymc(
+                inputs=[pytensor.In(v, borrow=True) for v in model.value_vars],
+                outputs=[pytensor.Out(v, borrow=True) for v in vars],
+                on_unused_input="ignore",
+            )
+            fn.trust_input = True
 
         # Get variable shapes. Most backends will need this
         # information.

pymc/backends/ndarray.py

@@ -76,7 +76,7 @@ def setup(self, draws, chain, sampler_vars=None) -> None:
         else:  # Otherwise, make array of zeros for each variable.
             self.draws = draws
             for varname, shape in self.var_shapes.items():
-                self.samples[varname] = np.zeros((draws, *shape), dtype=self.var_dtypes[varname])
+                self.samples[varname] = np.empty((draws, *shape), dtype=self.var_dtypes[varname])
 
         if sampler_vars is None:
             return
@@ -105,7 +105,7 @@ def record(self, point, sampler_stats=None) -> None:
         point: dict
             Values mapped to variable names
         """
-        for varname, value in zip(self.varnames, self.fn(point)):
+        for varname, value in zip(self.varnames, self.fn(*point.values())):
             self.samples[varname][self.draw_idx] = value
 
         if self._stats is not None and sampler_stats is None:

pymc/pytensorf.py

@@ -1024,7 +1024,12 @@ def compile_pymc(
     """
     # Create an update mapping of RandomVariable's RNG so that it is automatically
     # updated after every function call
-    rng_updates = collect_default_updates(inputs=inputs, outputs=outputs)
+    rng_updates = collect_default_updates(
+        inputs=[inp.variable if isinstance(inp, pytensor.In) else inp for inp in inputs],
+        outputs=[
+            out.variable if isinstance(out, pytensor.Out) else out for out in makeiter(outputs)
+        ],
+    )
 
     # We always reseed random variables as this provides RNGs with no chances of collision
     if rng_updates:

pymc/variational/opvi.py

@@ -1554,7 +1554,10 @@ def sample(
         if random_seed is not None:
             (random_seed,) = _get_seeds_per_chain(random_seed, 1)
         samples: dict = self.sample_dict_fn(draws, random_seed=random_seed)
-        points = ({name: records[i] for name, records in samples.items()} for i in range(draws))
+        points = (
+            {name: np.asarray(records[i]) for name, records in samples.items()}
+            for i in range(draws)
+        )
 
         trace = NDArray(
             model=self.model,

tests/backends/fixtures.py

@@ -238,7 +238,10 @@ class SamplingTestCase(ModelBackendSetupTestCase):
     """
 
     def record_point(self, val):
-        point = {varname: np.tile(val, value.shape) for varname, value in self.test_point.items()}
+        point = {
+            varname: np.tile(val, value.shape).astype(value.dtype)
+            for varname, value in self.test_point.items()
+        }
         if self.sampler_vars is not None:
             stats = [{key: dtype(val) for key, dtype in vars.items()} for vars in self.sampler_vars]
             self.strace.record(point=point, sampler_stats=stats)