Predictive log likelihood doesn't work for models with unplated observations.

[ThomasHeap]

Something like this should work but doesn't:

import torch as t
import torch.nn as nn
import alan
t.manual_seed(0)
def P(tr):
  '''
  Bayesian Gaussian Model
  '''
  a = t.zeros(5)
  tr.sample('mu', alan.Normal(a, 1))
  tr.sample('obs', alan.Normal(tr['mu'], 1))

class Q(alan.QModule):
    def __init__(self):
        super().__init__()
        self.m_mu = nn.Parameter(t.zeros((5,)))
        self.log_s_mu  = nn.Parameter(t.zeros((5,)))

    def forward(self, tr):
        tr.sample('mu', alan.Normal(self.m_mu, self.log_s_mu.exp()))# plates="plate_1")

data = alan.sample(P, varnames=('obs',))
test_data = alan.sample(P, varnames=('obs',))
all_data = {'obs': t.vstack([data['obs'],test_data['obs']])}

model = alan.Model(P, Q(), {'obs': data['obs']})

K = 10
pred_lik = model.predictive_ll(K, 1000, data_all=all_data)

We get the following error message:

Traceback (most recent call last):
  File "/home/thomas/Work/alan/examples/testing_predll.py", line 33, in <module>
    pred_lik = model.predictive_ll(K, 1000, data_all=all_data)
  File "/home/thomas/Work/alan/alan/model.py", line 290, in predictive_ll
    trace_pred, N = self._predictive(K, N, data_all, covariates_all, None)
  File "/home/thomas/Work/alan/alan/model.py", line 258, in _predictive
    self.P(tr)
  File "/home/thomas/Work/alan/examples/testing_predll.py", line 11, in P
    tr.sample('obs', alan.Normal(tr['mu'], 1))
  File "/home/thomas/Work/alan/alan/traces.py", line 558, in sample
    self._sample_logp(varname, dist, plates)
  File "/home/thomas/Work/alan/alan/traces.py", line 622, in _sample_logp
    ll_all                 = dist.log_prob(sample_all)
  File "/home/thomas/Work/alan/alan/dist.py", line 114, in log_prob
    assert x.ndim == self.result_ndim + self.unnamed_batch_dims
AssertionError

I think this is because the appending the data creates a dimension which isn't accounted for in the dimensions of the parameters of the 'obs' dist.

[LaurenceA]

There's no plates in P.

[ThomasHeap]

So something like:

import torch as t
import torch.nn as nn
import alan

def P(tr):
  '''
  Bayesian Gaussian Model
  '''
  a = t.zeros((5,))
  tr('mu', alan.Normal(a, 1))
  tr('obs', alan.Normal(tr['mu'], 1))

class Q(alan.AlanModule):
    def __init__(self):
        super().__init__()
        self.m_mu = nn.Parameter(t.zeros((5,)))
        self.log_s_mu  = nn.Parameter(t.zeros((5,)))

    def forward(self, tr):
        tr('mu', alan.Normal(self.m_mu, self.log_s_mu.exp()))

model = alan.Model(P, Q())
data = model.sample_prior(varnames=('obs',))
test_data = model.sample_prior(varnames=('obs',))

all_data = {'obs': t.vstack([data['obs'],test_data['obs']])}
model = model.condition(data=data)
sample = model.sample_cat(5, True)

pred_ll = model.predictive_ll(sample, 10, data_all=all_data)
print(pred_ll)

Which gives:

Traceback (most recent call last):
  File "/home/thomas/Work/alan/small_test_no_plates.py", line 35, in <module>
    pred_ll = model.predictive_ll(sample, 10, data_all=all_data)
  File "/home/thomas/Work/alan/alan/model.py", line 143, in predictive_ll
    trace_pred, N = self._predictive(sample, N, data_all, inputs_all, None)
  File "/home/thomas/Work/alan/alan/model.py", line 116, in _predictive
    tr = traces.TracePred(
  File "/home/thomas/Work/alan/alan/traces.py", line 457, in __init__
    raise Exception(f"None of the data tensors or plate sizes provided for prediction is bigger than those at training time.  Remember that the data/plate sizes are the sizes of train + 'test'")
Exception: None of the data tensors or plate sizes provided for prediction is bigger than those at training time.  Remember that the data/plate sizes are the sizes of train + 'test'

Must have plating for predictive_ll to work?

If I add plating:

import torch as t
import torch.nn as nn
import alan

def P(tr):
  '''
  Bayesian Gaussian Model
  '''
  a = t.zeros((5,))
  tr('mu', alan.Normal(a, 1))
  tr('obs', alan.Normal(tr['mu'], 1), plates='plate_1')

class Q(alan.AlanModule):
    def __init__(self):
        super().__init__()
        self.m_mu = nn.Parameter(t.zeros((5,)))
        self.log_s_mu  = nn.Parameter(t.zeros((5,)))

    def forward(self, tr):
        tr('mu', alan.Normal(self.m_mu, self.log_s_mu.exp()))# plates="plate_1")

model = alan.Model(P, Q())
data = model.sample_prior(varnames=('obs',), platesizes={"plate_1": 2})
test_data = model.sample_prior(varnames=('obs',), platesizes={"plate_1": 2})

all_data = {'obs': t.vstack([data['obs'],test_data['obs']])}
model = model.condition(data=data)
sample = model.sample_cat(5, True)

pred_ll = model.predictive_ll(sample, 10, data_all=all_data)
print(pred_ll)

I get the following error:

Traceback (most recent call last):
  File "/home/thomas/Work/alan/small_test.py", line 32, in <module>
    pred_ll = model.predictive_ll(sample, 10, data_all=all_data)
  File "/home/thomas/Work/alan/alan/model.py", line 143, in predictive_ll
    trace_pred, N = self._predictive(sample, N, data_all, inputs_all, None)
  File "/home/thomas/Work/alan/alan/model.py", line 123, in _predictive
    self.P(tr, **inputs_all)
  File "/home/thomas/Work/alan/alan/model.py", line 246, in P
    self.model.P(tr, *args, **kwargs)
  File "/home/thomas/Work/alan/small_test.py", line 11, in P
    tr('obs', alan.Normal(tr['mu'], 1), plates='plate_1')
  File "/home/thomas/Work/alan/alan/traces.py", line 59, in __call__
    self.sample_(key, dist, plates=plates, T=T, **kwargs)
  File "/home/thomas/Work/alan/alan/traces.py", line 490, in sample_
    self._sample_logp(varname, dist, plates)
  File "/home/thomas/Work/alan/alan/traces.py", line 554, in _sample_logp
    print(generic_order(ll_all, dims_all)[idxs])
ValueError: at least 2 indices were supplied but the tensor only has 1 dimensions

[LaurenceA]

Whenever we do "prediction" the variables we're predicting are all extensions of a plate. We can't "extend" something that isn't a plate...

On Thu, Mar 9, 2023 at 10:56 AM thomas @.***> wrote:

So something like:

import torch as timport torch.nn as nnimport alan def P(tr): ''' Bayesian Gaussian Model ''' a = t.zeros((5,)) tr('mu', alan.Normal(a, 1)) tr('obs', alan.Normal(tr['mu'], 1))

class Q(alan.AlanModule): def init(self): super().init() self.m_mu = nn.Parameter(t.zeros((5,))) self.log_s_mu = nn.Parameter(t.zeros((5,)))

def forward(self, tr):
    tr('mu', alan.Normal(self.m_mu, self.log_s_mu.exp()))

model = alan.Model(P, Q())data = model.sample_prior(varnames=('obs',))test_data = model.sample_prior(varnames=('obs',)) all_data = {'obs': t.vstack([data['obs'],test_data['obs']])}model = model.condition(data=data)sample = model.sample_cat(5, True) pred_ll = model.predictive_ll(sample, 10, data_all=all_data)print(pred_ll)

Which gives:

Traceback (most recent call last): File "/home/thomas/Work/alan/small_test_no_plates.py", line 35, in pred_ll = model.predictive_ll(sample, 10, data_all=all_data) File "/home/thomas/Work/alan/alan/model.py", line 143, in predictive_ll trace_pred, N = self._predictive(sample, N, data_all, inputs_all, None) File "/home/thomas/Work/alan/alan/model.py", line 116, in _predictive tr = traces.TracePred( File "/home/thomas/Work/alan/alan/traces.py", line 457, in init raise Exception(f"None of the data tensors or plate sizes provided for prediction is bigger than those at training time. Remember that the data/plate sizes are the sizes of train + 'test'") Exception: None of the data tensors or plate sizes provided for prediction is bigger than those at training time. Remember that the data/plate sizes are the sizes of train + 'test'

Must have plating for predictive_ll to work?

If I add plating:

import torch as timport torch.nn as nnimport alan def P(tr): ''' Bayesian Gaussian Model ''' a = t.zeros((5,)) tr('mu', alan.Normal(a, 1)) tr('obs', alan.Normal(tr['mu'], 1), plates='plate_1')

class Q(alan.AlanModule): def init(self): super().init() self.m_mu = nn.Parameter(t.zeros((5,))) self.log_s_mu = nn.Parameter(t.zeros((5,)))

def forward(self, tr):
    tr('mu', alan.Normal(self.m_mu, self.log_s_mu.exp()))# plates="plate_1")

model = alan.Model(P, Q())data = model.sample_prior(varnames=('obs',), platesizes={"plate_1": 2})test_data = model.sample_prior(varnames=('obs',), platesizes={"plate_1": 2}) all_data = {'obs': t.vstack([data['obs'],test_data['obs']])}model = model.condition(data=data)sample = model.sample_cat(5, True) pred_ll = model.predictive_ll(sample, 10, data_all=all_data)print(pred_ll)

I get the following error:

Traceback (most recent call last): File "/home/thomas/Work/alan/small_test.py", line 32, in pred_ll = model.predictive_ll(sample, 10, data_all=all_data) File "/home/thomas/Work/alan/alan/model.py", line 143, in predictive_ll trace_pred, N = self._predictive(sample, N, data_all, inputs_all, None) File "/home/thomas/Work/alan/alan/model.py", line 123, in _predictive self.P(tr, *inputs_all) File "/home/thomas/Work/alan/alan/model.py", line 246, in P self.model.P(tr, args, kwargs) File "/home/thomas/Work/alan/small_test.py", line 11, in P tr('obs', alan.Normal(tr['mu'], 1), plates='plate1') File "/home/thomas/Work/alan/alan/traces.py", line 59, in call self.sample(key, dist, plates=plates, T=T, kwargs) File "/home/thomas/Work/alan/alan/traces.py", line 490, in sample_ self._sample_logp(varname, dist, plates) File "/home/thomas/Work/alan/alan/traces.py", line 554, in _sample_logp print(generic_order(ll_all, dims_all)[idxs]) ValueError: at least 2 indices were supplied but the tensor only has 1 dimensions

— Reply to this email directly, view it on GitHub https://github.com/alan-ppl/alan/issues/2#issuecomment-1461801633, or unsubscribe https://github.com/notifications/unsubscribe-auth/ABAUSQV5ADHZI7KIVZ2HRK3W3GZOPANCNFSM6AAAAAAT2O7VXU . You are receiving this because you commented.Message ID: @.***>

[ThomasHeap]

Ok, so in the first case that's the intended behaviour and I'll make a note of it in the documentation as I expect this might trip new users up.

But there is still the issue with the second example:

import torch as t
import torch.nn as nn
import alan

def P(tr):
  '''
  Bayesian Gaussian Model
  '''
  a = t.zeros((5,))
  tr('mu', alan.Normal(a, 1))
  tr('obs', alan.Normal(tr['mu'], 1), plates='plate_1')

class Q(alan.AlanModule):
    def __init__(self):
        super().__init__()
        self.m_mu = nn.Parameter(t.zeros((5,)))
        self.log_s_mu  = nn.Parameter(t.zeros((5,)))

    def forward(self, tr):
        tr('mu', alan.Normal(self.m_mu, self.log_s_mu.exp()), plates="plate_1")

model = alan.Model(P, Q())
data = model.sample_prior(varnames=('obs',), platesizes={"plate_1": 2})
test_data = model.sample_prior(varnames=('obs',), platesizes={"plate_1": 2})

all_data = {'obs': t.vstack([data['obs'],test_data['obs']])}
model = model.condition(data=data)
sample = model.sample_cat(5, True)

pred_ll = model.predictive_ll(sample, 10, data_all=all_data)
print(pred_ll)

which returns

Traceback (most recent call last):
  File "/home/thomas/Work/alan/small_test.py", line 32, in <module>
    pred_ll = model.predictive_ll(sample, 10, data_all=all_data)
  File "/home/thomas/Work/alan/alan/model.py", line 144, in predictive_ll
    trace_pred, N = self._predictive(sample, N, data_all, inputs_all, None)
  File "/home/thomas/Work/alan/alan/model.py", line 124, in _predictive
    self.P(tr, **inputs_all)
  File "/home/thomas/Work/alan/alan/model.py", line 247, in P
    self.model.P(tr, *args, **kwargs)
  File "/home/thomas/Work/alan/small_test.py", line 10, in P
    tr('mu', alan.Normal(a, 1))
  File "/home/thomas/Work/alan/alan/traces.py", line 59, in __call__
    self.sample_(key, dist, plates=plates, T=T, **kwargs)
  File "/home/thomas/Work/alan/alan/traces.py", line 496, in sample_
    self._sample_sample(varname, dist, plates)
  File "/home/thomas/Work/alan/alan/traces.py", line 508, in _sample_sample
    sample       = generic_order(sample,       dims_all)   #Still torchdim, as it has N!
  File "/home/thomas/Work/alan/alan/utils.py", line 68, in generic_order
    assert_no_ellipsis(dims)
  File "/home/thomas/Work/alan/alan/utils.py", line 78, in assert_no_ellipsis
    assert dims[-1] is not Ellipsis
AssertionError

I think is because self.corresponding_plates adds an Ellipsis dimension to the end of dims_all.

[LaurenceA]

Should be fixed.

[ThomasHeap]

With

import torch as t
import torch.nn as nn
import alan

def P(tr):
  '''
  Bayesian Gaussian Model
  '''
  a = t.zeros((5,))
  tr('mu', alan.Normal(a, 1))
  tr('obs', alan.Normal(tr['mu'], 1), plates='plate_1')

class Q(alan.AlanModule):
    def __init__(self):
        super().__init__()
        self.m_mu = nn.Parameter(t.zeros((5,)))
        self.log_s_mu  = nn.Parameter(t.zeros((5,)))

    def forward(self, tr):
        tr('mu', alan.Normal(self.m_mu, self.log_s_mu.exp()))# plates="plate_1")

model = alan.Model(P, Q())
data = model.sample_prior(varnames=('obs',), platesizes={"plate_1": 2})
test_data = model.sample_prior(varnames=('obs',), platesizes={"plate_1": 2})

all_data = {'obs': t.vstack([data['obs'],test_data['obs']])}
model = model.condition(data=data)
sample = model.sample_cat(5, True)

pred_ll = model.predictive_ll(sample, 10, data_all=all_data)
print(pred_ll)

I get the following error:

Traceback (most recent call last):
  File "/home/thomas/Work/alan/small_test.py", line 30, in <module>
    sample = model.sample_cat(5, True)
  File "/home/thomas/Work/alan/alan/model.py", line 66, in sample_cat
    return self.sample_tensor(traces.TraceQCategorical, *args, **kwargs)
  File "/home/thomas/Work/alan/alan/model.py", line 100, in sample_tensor
    platedims, data, inputs = self.dims_data_inputs(data, inputs, platesizes, device)
  File "/home/thomas/Work/alan/alan/model.py", line 51, in dims_data_inputs
    data   = named2dim_tensordict(platedims, data)
  File "/home/thomas/Work/alan/alan/utils.py", line 224, in named2dim_tensordict
    return {k: named2dim_tensor(d, tensor) for (k, tensor) in tensordict.items()}
  File "/home/thomas/Work/alan/alan/utils.py", line 224, in <dictcomp>
    return {k: named2dim_tensor(d, tensor) for (k, tensor) in tensordict.items()}
  File "/home/thomas/Work/alan/alan/utils.py", line 210, in named2dim_tensor
    raise Exception(f"No torchdim dimension for named dimension {name} in named2dim_tensor")
Exception: No torchdim dimension for named dimension None in named2dim_tensor

Can we fix this by replacing:

https://github.com/alan-ppl/alan/blob/275702a98adde3cb61f01aa29e65df45add98e78/alan/utils.py#L209

with:

if name not in d and name is not None: