Memory Error: OrderedLogistic

[aloctavodia]

I am updating and re-running the statistical rethinking notebooks. I get a memory allocation error with model m_11 (code block 11.5). The problem seems to be related with #3383, reverting changes in categorical distribution previous to that PR, fix the issue. @lucianopaz probably you have a better idea of what is going on here.

[tohtsky]

I think the basic example provided in OrderedLogistic (slightly modified to use ADVI) behaves differently in the recent commits.
I've run the following code in the current head and old commits like 9ef2947.

n1_c = 300; n2_c = 300; n3_c = 300
cluster1 = np.random.randn(n1_c) + -1
cluster2 = np.random.randn(n2_c) + 0
cluster3 = np.random.randn(n3_c) + 2

x = np.concatenate((cluster1, cluster2, cluster3))
y = np.concatenate((1*np.ones(n1_c),
                    2*np.ones(n2_c),
                    3*np.ones(n3_c))) - 1

# Ordered logistic regression
with pm.Model() as model:
    cutpoints = pm.Normal("cutpoints", mu=[-1,1], sigma=10, shape=2,
                          transform=pm.distributions.transforms.ordered)
    y_ = pm.OrderedLogistic("y", cutpoints=cutpoints, eta=x, observed=y)
    inference = pm.fit(30000, method='advi')
    tr = inference.sample(1000)

In the older versions, the value of ELBO seems to converge to a reasonable value ~ 10^2, but in the recent commits, the progress bar shows large value ~ 10 ^5.

I suspect this is caused by logp method in Categorical.
In the older versions, when p.ndim == 2 (which is the case in the above example), the log likelihood

a = tt.log(p[tt.arange(p.shape[0]), value_clip])

will be a length p.shape[0] vector.
In recent versions, however, a defined as

pattern = (p.ndim - 1,) + tuple(range(p.ndim - 1))
a = tt.log(p.dimshuffle(pattern)[value_clip])

will be a p.shape[0] x p.shape[0] matrix (which has 900 **2 entries in the above example, roughly matching with the scale of ELBO).

[junpenglao]

Thanks @tohtsky, I think your diagnosis is right.

[lucianopaz]

@aloctavodia, sorry for the delay in taking a look at this. I'm still mostly offline until next week.

I looked a bit more into the commit you referenced and I think that there may be a bug in Categorical.logp, specifically here. If p is multidimensional, the dimshuffle moves the last axis around to simplify advanced indexing into it with the supplied values. The resulting a will have a weird shape like (value.shape, p.shape[:-1]), instead of having value.shape at the end. This could lead to strange broadcasting with observeds and produce an intermediate array of shape (value.shape, value.shape). So this could cause the memory error you encountered. Maybe that line could be replaced with this:

a = tt.log(p[..., value_clip])

I'll be able to test this in a bit less than two weeks, but if you want to try it out first, you're welcome to do so.

[lucianopaz]

Oops, I skipped over @tohtsky's answer that points to exactly the same line of code that I thought was causing the bug in my previous answer. Maybe we could try with the ellipses in the indexing instead of doing a dimshuffle.

[aloctavodia]

I tried replacing the line you suggested and I still see the same problem :-(

[tohtsky]

Simply flattening multi-dimensional tensor into a 2d array and then reshaping the resulting logp vector back into the original tensor shape should work?

        if p.ndim > 1:
            original_shape = p.shape[:-1]
            p_flatten = p.reshape((-1, p.shape[-1]))
            a = tt.log(
                p_flatten[
                    (tt.arange(p_flatten.shape[0]), value_clip.ravel())
                ]).reshape(
                    original_shape
            )
        else:
            a = tt.log(p[value_clip])

[lucianopaz]

I'll have time to look into this issue this Friday.

[bdyetton]

Hi,
Any progress on this issue?

[lucianopaz]

I managed to write a small test involving logp that highlights the error:

import numpy as np
from scipy.special import logit
import pymc3 as pm


loge = np.log10(np.exp(1))
size = 100
p = np.ones(10) / 10
cutpoints = logit(np.linspace(0, 1, 11)[1:-1])
obs = np.random.randint(0, 1, size=size)
with pm.Model():
    ol = pm.OrderedLogistic("ol", eta=0, cutpoints=cutpoints, observed=obs)
    c = pm.Categorical("c", p=p, observed=obs)

print(ol.logp({"ol": 1}) * loge)
print(c.logp({"c": 1}) * loge)

The OrderedLogistic variable ol, given the provided cutpoints, should be equivalent to the categorical RV c. When we look at the returned logp for each RV on a given class, as all are equally likely, the log10 should be approximately equal to -size == -100. This is the case for the Categorical but the OrderedLogistic returns -size**2 == -10000. The problem seems to be in an unnecessary shape padding done in the OrderedLogistic.__init__.

Furthermore,

>>> ol.distribution.p.ndim
2
>>> c.distribution.p.ndim
1

However, there is also an additional shape issue when we give an array of cutpoints, and how this broadcasts with the passed observed. I'll try to finish a fix tomorrow.

[lucianopaz]

In the end, it wasn't a broadcasting problem. It was an indexing problem when p was multidimensional.