Logistic regression with time-varying coefficients: cannot construct sampler about aemcmc HOT 6 OPEN

import numpy as np

import aesara
import aesara.tensor as at
from aemcmc.basic import construct_sampler
from aesara.tensor.random.utils import RandomStream

def logistic_fit(X_val, y_val):
    N, M, T = X_val.shape

    srng = RandomStream(0)
    X = at.tensor3("X")

    sigma_rv = srng.exponential(1, size=X.shape[1])
    beta_t_rv = at.cumsum(srng.normal(0, 1/sigma_rv, size=(X.shape[1],X.shape[2])), axis=1)

    eta = at.tensordot(X, beta_t_rv, 2)
    p = at.sigmoid(-eta)
    Y_rv = srng.bernoulli(p, name="Y")

    y_vv = Y_rv.clone()
    y_vv.name = "y"

    sample_vars = [sigma_rv, beta_t_rv]

    sampler, initial_values = construct_sampler({Y_rv: y_vv}, srng)

    inputs = [X, y_vv] + [initial_values[rv] for rv in sample_vars]
    outputs = [sampler.sample_steps[rv] for rv in sample_vars]

    sample_step = aesara.function(
        inputs,
        outputs,
        updates=sampler.updates,
        on_unused_input="ignore",
    )

    sigma_val = np.ones(M)

    beta_pst_vals = []

    sigma_pst_val, beta_pst_val = (
        sigma_val,
        np.zeros(M,T)
    )
    for i in range(100):
        sigma_pst_val, beta_pst_val = sample_step(
            X_val,
            y_val,
            sigma_pst_val,
            beta_pst_val
        )
        beta_pst_vals += [beta_pst_val]
    
    beta_pst_mean = np.mean(beta_pst_vals, axis=0)
    return beta_pst_mean

# X_val = np.load("X_val.npy")
# y_val = np.load("y_val.npy")

X_val = np.zeros((1000, 50, 10))
y_val = np.zeros(1000)

beta = logistic_fit(X_val, y_val)

/Users/acristia/anaconda3/lib/python3.8/site-packages/aehmc/utils.py:43: UserWarning: The following parameters need to be computed in order to determine the shapes in this parameter map: [<TensorType(float64, (None, None))>]
  warnings.warn(
Traceback (most recent call last):
  File "examples/gibbs_sample.py", line 61, in <module>
    beta = logistic_fit(X_val, y_val)
  File "examples/gibbs_sample.py", line 28, in logistic_fit
    inputs = [X, y_vv] + [initial_values[rv] for rv in sample_vars]
  File "examples/gibbs_sample.py", line 28, in <listcomp>
    inputs = [X, y_vv] + [initial_values[rv] for rv in sample_vars]
KeyError: CumOp{1, add}.0