Logistic regression with time-varying coefficients: cannot construct sampler

[lucasgautheron]

Description of your problem or feature request

I want to perform a classification task which consists in predicting whether a document $i$ published at a time $t_i$ belongs to a certain class ( $y_i=1$ ) or not ( $y_i=0$ ). The contents of $i$ is given by a bag of words $b_{ij}$. The predictive features are encoded into a sparse tensor $x_{ijt} = b_{ij} \delta_{t,t_i}$. For this task I am considering a logistic regression with time-varying coefficients, such that the coefficients evolve through a random walk:

$$\text{logit} (p_{i}) = \sum_{jt} x_{ijt} \beta_{jt}$$

$$y_{i} \sim \text{Bernouilli}(p_i)$$

$$\beta_{j,t+1} \sim \text{Normal}(\beta_{j,t},\sigma_j^{-1})$$

$$\beta_{j,0} \sim \text{Normal}(0,\sigma_{j,0}^{-1})$$

$$\sigma_j \sim \text{Exponential}(1)$$

$$\sigma_{j,0} \sim \text{Exponential}(1)$$

For the sake of simplicity below I assume $\sigma_{j,0}=\sigma_j$, although I do not see any compelling reason to do so.

However:

• I cannot construct a sampler for this model (see details below)
• I wonder if there could be a more efficient implementation anyway: $x_{ijt}$ is very sparse, and at.tensordot(X, beta_t_rv, 2) involves too many needless operations. I would not have done that with Stan but in this case I found it easier to write the model this way.

Please provide a minimal, self-contained, and reproducible example.

As a clueless user, I tried to build a sampler based on my generative model by looking at other examples; which means I could be doing something completely idiotic!
Here is what I came up with:

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]

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

# 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)

Although thus far the script does not actually use the data, if at some point you want to try with actual data, download and unzip these files (X_val.npy and y_val.npy) to the folder the script is executed from: data.zip

However, this crashes:

Please provide the full traceback of any errors.

ERROR (aesara.graph.rewriting.basic): Rewrite failure due to: local_elemwise_dimshuffle_subsume
ERROR (aesara.graph.rewriting.basic): node: Elemwise{true_div,no_inplace}(InplaceDimShuffle{x}.0, exponential_rv{0, (0,), floatX, False}.out)
ERROR (aesara.graph.rewriting.basic): TRACEBACK:
ERROR (aesara.graph.rewriting.basic): Traceback (most recent call last):
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aesara/graph/rewriting/basic.py", line 1933, in process_node
    replacements = node_rewriter.transform(fgraph, node)
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aesara/graph/rewriting/basic.py", line 1092, in transform
    return self.fn(fgraph, node)
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aemcmc/rewriting.py", line 345, in local_elemwise_dimshuffle_subsume
    new_op = SubsumingElemwise(new_inputs, new_outputs, inline=True)
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aemcmc/rewriting.py", line 189, in __init__
    OpFromGraph.__init__(self, inputs, outputs, *args, **kwargs)
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aesara/compile/builders.py", line 343, in __init__
    raise TypeError(f"Constants not allowed as inputs; {i}")
TypeError: Constants not allowed as inputs; TensorConstant{1}

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 26, in logistic_fit
    sample_steps, updates, initial_values = construct_sampler({Y_rv: y_vv}, srng)
  File "/Users/acristia/anaconda3/lib/python3.8/site-packages/aemcmc/basic.py", line 108, in construct_sampler
    raise NotImplementedError(
NotImplementedError: Could not find a posterior samplers for {exponential_rv{0, (0,), floatX, False}.out, normal_rv{0, (0, 0), floatX, False}.out}

Versions and main components

• Aesara version: 2.8.9
• Aemcmc version: 0.06
• Python version: 3.8.10
• Operating system: macOS 10.15.4
• How did you install Aesara: pip

[brandonwillard]

Thank for the bug report! We'll take a closer look at it shortly.

[brandonwillard]

I have a fix in #93 for the error you were seeing, but we still need to look into the samplers AeMCMC produces (or lack thereof) for your model.

[lucasgautheron]

Thank you, that was quick! I've tried with the patch, and construct_sampler no longer fails. However there is a new error, which could be very well be my mistake?

Here is the new code:

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)

Here is the error (also notice the warning)

/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

Let me know if I can provide more useful information!

[brandonwillard]

That CumOp warning is something we still need to address. We can open another issue for it, though.

[rlouf]

Looked like an error to me?

[brandonwillard]

Looked like an error to me?

Yes, it is. I'll reopen this and take a look soon.