Cache size in DictToArrayBijection

Author: ricardoV94

pymc/blocking.py

@@ -39,11 +39,11 @@
 StatShape: TypeAlias = Sequence[int | None] | None
 
 
-# `point_map_info` is a tuple of tuples containing `(name, shape, dtype)` for
+# `point_map_info` is a tuple of tuples containing `(name, shape, size, dtype)` for
 # each of the raveled variables.
 class RaveledVars(NamedTuple):
     data: np.ndarray
-    point_map_info: tuple[tuple[str, tuple[int, ...], np.dtype], ...]
+    point_map_info: tuple[tuple[str, tuple[int, ...], int, np.dtype], ...]
 
 
 class Compose(Generic[T]):
@@ -67,10 +67,9 @@ class DictToArrayBijection:
     @staticmethod
     def map(var_dict: PointType) -> RaveledVars:
         """Map a dictionary of names and variables to a concatenated 1D array space."""
-        vars_info = tuple((v, k, v.shape, v.dtype) for k, v in var_dict.items())
-        raveled_vars = [v[0].ravel() for v in vars_info]
-        if raveled_vars:
-            result = np.concatenate(raveled_vars)
+        vars_info = tuple((v, k, v.shape, v.size, v.dtype) for k, v in var_dict.items())
+        if vars_info:
+            result = np.concatenate(tuple(v[0].ravel() for v in vars_info))
         else:
             result = np.array([])
         return RaveledVars(result, tuple(v[1:] for v in vars_info))
@@ -91,19 +90,15 @@ def rmap(
 
         """
         if start_point:
-            result = dict(start_point)
+            result = start_point.copy()
         else:
             result = {}
 
-        if not isinstance(array, RaveledVars):
-            raise TypeError("`array` must be a `RaveledVars` type")
-
         last_idx = 0
-        for name, shape, dtype in array.point_map_info:
-            arr_len = np.prod(shape, dtype=int)
-            var = array.data[last_idx : last_idx + arr_len].reshape(shape).astype(dtype)
-            result[name] = var
-            last_idx += arr_len
+        for name, shape, size, dtype in array.point_map_info:
+            end = last_idx + size
+            result[name] = array.data[last_idx:end].reshape(shape).astype(dtype)
+            last_idx = end
 
         return result
 

pymc/sampling/parallel.py

@@ -228,15 +228,12 @@ def __init__(
         self._shared_point = {}
         self._point = {}
 
-        for name, shape, dtype in DictToArrayBijection.map(start).point_map_info:
-            size = 1
-            for dim in shape:
-                size *= int(dim)
-            size *= dtype.itemsize
-            if size != ctypes.c_size_t(size).value:
+        for name, shape, size, dtype in DictToArrayBijection.map(start).point_map_info:
+            byte_size = size * dtype.itemsize
+            if byte_size != ctypes.c_size_t(byte_size).value:
                 raise ValueError(f"Variable {name} is too large")
 
-            array = mp_ctx.RawArray("c", size)
+            array = mp_ctx.RawArray("c", byte_size)
             self._shared_point[name] = (array, shape, dtype)
             array_np = np.frombuffer(array, dtype).reshape(shape)
             array_np[...] = start[name]

pymc/step_methods/hmc/quadpotential.py

@@ -363,23 +363,23 @@ def raise_ok(self, map_info):
         if np.any(self._stds == 0):
             errmsg = ["Mass matrix contains zeros on the diagonal. "]
             last_idx = 0
-            for name, shape, dtype in map_info:
-                arr_len = np.prod(shape, dtype=int)
-                index = np.where(self._stds[last_idx : last_idx + arr_len] == 0)[0]
+            for name, shape, size, dtype in map_info:
+                end = last_idx + size
+                index = np.where(self._stds[last_idx:end] == 0)[0]
                 errmsg.append(f"The derivative of RV `{name}`.ravel()[{index}] is zero.")
-                last_idx += arr_len
+                last_idx += end
 
             raise ValueError("\n".join(errmsg))
 
         if np.any(~np.isfinite(self._stds)):
             errmsg = ["Mass matrix contains non-finite values on the diagonal. "]
 
             last_idx = 0
-            for name, shape, dtype in map_info:
-                arr_len = np.prod(shape, dtype=int)
-                index = np.where(~np.isfinite(self._stds[last_idx : last_idx + arr_len]))[0]
+            for name, shape, size, dtype in map_info:
+                end = last_idx + size
+                index = np.where(~np.isfinite(self._stds[last_idx:end]))[0]
                 errmsg.append(f"The derivative of RV `{name}`.ravel()[{index}] is non-finite.")
-                last_idx += arr_len
+                last_idx = end
             raise ValueError("\n".join(errmsg))
 
 

pymc/tuning/starting.py

@@ -143,7 +143,7 @@ def find_MAP(
     compiled_logp_func = DictToArrayBijection.mapf(model.compile_logp(jacobian=False), start)
     logp_func = lambda x: compiled_logp_func(RaveledVars(x, x0.point_map_info))  # noqa: E731
 
-    rvs = [model.values_to_rvs[vars_dict[name]] for name, _, _ in x0.point_map_info]
+    rvs = [model.values_to_rvs[vars_dict[name]] for name, _, _, _ in x0.point_map_info]
     try:
         # This might be needed for calls to `dlogp_func`
         # start_map_info = tuple((v.name, v.shape, v.dtype) for v in vars)