just use scan bro

Author: jessegrabowski

pymc/distributions/transforms.py

@@ -16,6 +16,7 @@
 from functools import singledispatch
 
 import numpy as np
+import pytensor
 import pytensor.tensor as pt
 
 
@@ -179,6 +180,33 @@ def __init__(self, n):
         """
         self.n = n
 
+    def step(self, i, counter, L, y):
+        y_star = y[counter : counter + i]
+        dsy = y_star.dot(y_star)
+        alpha_r = 1 / (dsy + 1)
+        gamma = pt.sqrt(dsy + 2) * alpha_r
+
+        x = pt.join(0, gamma * y_star, pt.atleast_1d(alpha_r))
+        next_L = L[i, : i + 1].set(x)
+        log_det = pt.log(2) + 0.5 * (i - 2) * pt.log(dsy + 2) - i * pt.log(1 + dsy)
+
+        return next_L, log_det
+
+    def _compute_L_and_logdet_scan(self, value, *inputs):
+        L = pt.eye(self.n)
+        idxs = pt.arange(1, self.n)
+        counters = pt.arange(0, self.n).cumsum()
+
+        results, _ = pytensor.scan(
+            self.step, outputs_info=[L, None], sequences=[idxs, counters], non_sequences=[value]
+        )
+
+        L_seq, log_det_seq = results
+        L = L_seq[-1]
+        log_det = pt.sum(log_det_seq)
+
+        return L, log_det
+
     def _compute_L_and_logdet(self, value, *inputs):
         n = self.n
         counter = 0
@@ -201,7 +229,7 @@ def _compute_L_and_logdet(self, value, *inputs):
         return L, log_det
 
     def backward(self, value, *inputs):
-        L, _ = self._compute_L_and_logdet(value, *inputs)
+        L, _ = self._compute_L_and_logdet_scan(value, *inputs)
         return L
 
     def forward(self, value, *inputs):
@@ -211,7 +239,7 @@ def forward(self, value, *inputs):
         return pt.as_tensor_variable(np.random.normal(size=size))
 
     def log_jac_det(self, value, *inputs):
-        _, log_det = self._compute_L_and_logdet(value, *inputs)
+        _, log_det = self._compute_L_and_logdet_scan(value, *inputs)
         return log_det