Sample within theano.scan?

Questions
1

Is it possible to sample within a loop? This piece of code works fine:

n_subj = 5
n_trials = 4
n_actions = 3

with pm.Model() as model:
    p = pm.Uniform('p', lower=0, upper=1, shape=n_subj, testval=0.1 + 0.1 * np.arange(n_subj))
    p_tile = T.repeat(T.tile(p, n_trials), n_actions).reshape([n_trials, n_subj, n_actions])

    def fun(p_tile):
        # Function containing sampling
        choice = pm.Categorical('choice', p=p_tile, shape=n_subj, testval=np.random.choice(range(n_actions), n_subj))
        return choice

    # Call sampling function once
    choices = fun(p_tile[0])  
    trace = pm.sample(500, tune=500, chains=1, cores=1)

But this one crashes:

n_subj = 5
n_trials = 4
n_actions = 3

with pm.Model() as model:
    p = pm.Uniform('p', lower=0, upper=1, shape=n_subj, testval=0.1 + 0.1 * np.arange(n_subj))
    p_tile = T.repeat(T.tile(p, n_trials), n_actions).reshape([n_trials, n_subj, n_actions])

    def fun(p_tile):
        # Function containing sampling
        choice = pm.Categorical('choice', p=p_tile, shape=n_subj, testval=np.random.choice(range(n_actions), n_subj))
        return choice

    # Call sampling function inside a loop
    choices, _ = theano.scan(fn=fun,  
                             sequences=[p_tile])
    trace = pm.sample(500, tune=500, chains=1, cores=1)

Here is the entire error message:

Traceback (most recent call last):
  File "test.py", line 31, in <module>
    trace = pm.sample(500, tune=500, chains=1, cores=1)
  File "C:\Users\maria\Anaconda3\lib\site-packages\pymc3\sampling.py", line 405, in sample
    step = assign_step_methods(model, step, step_kwargs=step_kwargs)
  File "C:\Users\maria\Anaconda3\lib\site-packages\pymc3\sampling.py", line 149, in assign_step_methods
    return instantiate_steppers(model, steps, selected_steps, step_kwargs)
  File "C:\Users\maria\Anaconda3\lib\site-packages\pymc3\sampling.py", line 70, in instantiate_steppers
    step = step_class(vars=vars, **args)
  File "C:\Users\maria\Anaconda3\lib\site-packages\pymc3\step_methods\hmc\nuts.py", line 152, in __init__
    super(NUTS, self).__init__(vars, **kwargs)
  File "C:\Users\maria\Anaconda3\lib\site-packages\pymc3\step_methods\hmc\base_hmc.py", line 63, in __init__
    dtype=dtype, **theano_kwargs)
  File "C:\Users\maria\Anaconda3\lib\site-packages\pymc3\step_methods\arraystep.py", line 215, in __init__
    vars, dtype=dtype, **theano_kwargs)
  File "C:\Users\maria\Anaconda3\lib\site-packages\pymc3\model.py", line 708, in logp_dlogp_function
    return ValueGradFunction(self.logpt, grad_vars, extra_vars, **kwargs)
  File "C:\Users\maria\Anaconda3\lib\site-packages\pymc3\model.py", line 447, in __init__
    inputs, [self._cost_joined, grad], givens=givens, **kwargs)
  File "C:\Users\maria\Anaconda3\lib\site-packages\theano\compile\function.py", line 317, in function
    output_keys=output_keys)
  File "C:\Users\maria\Anaconda3\lib\site-packages\theano\compile\pfunc.py", line 486, in pfunc
    output_keys=output_keys)
  File "C:\Users\maria\Anaconda3\lib\site-packages\theano\compile\function_module.py", line 1839, in orig_function
    name=name)
  File "C:\Users\maria\Anaconda3\lib\site-packages\theano\compile\function_module.py", line 1487, in __init__
    accept_inplace)
  File "C:\Users\maria\Anaconda3\lib\site-packages\theano\compile\function_module.py", line 181, in std_fgraph
    update_mapping=update_mapping)
  File "C:\Users\maria\Anaconda3\lib\site-packages\theano\gof\fg.py", line 175, in __init__
    self.__import_r__(output, reason="init")
  File "C:\Users\maria\Anaconda3\lib\site-packages\theano\gof\fg.py", line 346, in __import_r__
    self.__import__(variable.owner, reason=reason)
  File "C:\Users\maria\Anaconda3\lib\site-packages\theano\gof\fg.py", line 391, in __import__
    raise MissingInputError(error_msg, variable=r)
theano.gof.fg.MissingInputError: Input 0 of the graph (indices start from 0), used to compute Sum{axis=[1], acc_dtype=float64}(<TensorType(float64, matrix)>), was not provided and not given a value. Use the Theano flag exception_verbosity='high', for more information on this error.
 
Backtrace when that variable is created:

  File "test.py", line 26, in <module>
    sequences=[p_tile])


Process finished with exit code 1

In this example, I could just get rid of the loop, but I need the loop in my actual model. Is there any way to sample within a loop? Thanks in advance!!

(The actual model looks more like this:

# Get data
n_seasons, n_TS, n_aliens, n_actions = 3, 3, 4, 3
n_subj, n_trials = 2, 5
seasons = np.random.choice(range(n_seasons), size=[n_trials, n_subj])
aliens = np.random.choice(range(n_aliens), size=[n_trials, n_subj])
actions = np.random.choice(range(n_actions), size=[n_trials, n_subj])
rewards = 10 * np.random.rand(n_trials * n_subj).reshape([n_trials, n_subj]).round(2)
initial_Q = 1.6

# Convert data to tensor variables
seasons = T.as_tensor_variable(seasons)
aliens = T.as_tensor_variable(aliens)
actions = T.as_tensor_variable(actions)
rewards = T.as_tensor_variable(rewards)

trials, subj = np.meshgrid(range(n_trials), range(n_subj))
trials = T.as_tensor_variable(trials.T)
subj = T.as_tensor_variable(subj.T)

# Make RL model
with pm.Model() as model:
    alpha = pm.Beta('alpha', alpha=2, beta=2, shape=n_subj, testval=np.random.choice([0.1, 0.5], n_subj))
    beta = pm.Gamma('beta', alpha=8, beta=1, shape=n_subj, testval=5 * np.random.rand(n_subj).round(2))
    betas = T.tile(T.repeat(beta, n_actions), n_trials).reshape([n_trials, n_subj, n_actions])    # Q_sub.shape

    beta_high = beta.copy()
    beta_highs = T.repeat(beta_high, n_TS).reshape([n_subj, n_TS])  # Q_high_sub.shape

    forget = pm.Beta('forget', alpha=2, beta=2, shape=n_subj, testval=np.random.choice([0.1, 0.5], n_subj))
    forgets = T.repeat(forget, n_TS * n_aliens * n_actions).reshape([n_subj, n_TS, n_aliens, n_actions])  # Q_low for 1 trial

    # Initial Q-values
    Q_low0 = initial_Q * T.ones([n_subj, n_TS, n_aliens, n_actions])
    Q_high0 = initial_Q * T.ones([n_subj, n_seasons, n_TS])

    # Function to update Q-values based on stimulus, action, and reward
    def update_Q(season, alien, action, reward, Q_low, Q_high, alpha, beta_high, forget):
        # Loop over trials: take data for all subjects, 1 single trial
        Q_high_sub = Q_high[T.arange(n_subj), season] 
        p_high = T.exp(beta_high * Q_high_sub)

        # THIS LINE MAKES THE SCRIPT CRASH
        TS = pm.Categorical('TS', p=p_high, shape=n_subj, testval=np.ones(n_subj))
        # THIS LINE WORKS FINE
        # TS = T.as_tensor_variable(np.ones(n_subj).astype(int))
        T.printing.Print('TS')(TS)

        Q_low_new = (1 - forget) * Q_low
        RPE_low = alpha * (reward - Q_low_new[T.arange(n_subj), TS, alien, action])
        Q_low_new = T.set_subtensor(Q_low_new[T.arange(n_subj), TS, alien, action],
                                    Q_low_new[T.arange(n_subj), TS, alien, action] + RPE_low)
        return [Q_low_new, Q_high]

    # Get Q-values for the whole task (update each trial)
    [Q_low, Q_high], _ = theano.scan(fn=update_Q,
                                     sequences=[seasons, aliens, actions, rewards],
                                     outputs_info=[Q_low0, Q_high0],
                                     non_sequences=[alpha, beta_highs, forgets])

    Q_low = T.concatenate([[Q_low0], Q_low[:-1]], axis=0)  # Add first trial's Q-values, remove last trials Q-values

    # Select the right Q-values for each trial
    Q_sub = Q_low[trials, subj, seasons, aliens]

    # First step to translate into probabilities: apply exp
    p = T.exp(betas * Q_sub)

    # Bring p in the correct shape for pm.Categorical: 2-d array of trials
    action_wise_p = p.reshape([n_trials * n_subj, n_actions])
    action_wise_actions = actions.flatten()

    # Select actions (& second step of calculating probabilities: normalize)
    actions = pm.Categorical('actions', p=action_wise_p, observed=action_wise_actions)

    # Sample the model
    trace = pm.sample(100, tune=50, chains=1, cores=1)

This model throws the same error message as the piece of code above.)

2

I think it doesnt work if you are trying to create a pymc3 RandomVariable in the function you are performing scan. If I read it correctly you are trying to create a indexing tensor using Categorical distribution, but I think an easier way is either rewrite the scan with categorical into a distribution, or more simple do p_high.argmax() instead.

3

Thanks for the fast reply! That’s exactly what I’m trying to do.

I can’t use p_high.argmax() here because I need to sample according to p. How do I rewrite the scan loop into a distribution? Do I need to make a new theano op, or a custom pymc3 distribution, or something else? I can’t think of a way to get around the loop because I only get the probabilities for each Categorical sample from the probabilities of the previous sample.

4

Since you have a stochastic function, I would try to rewrite it as a custom distribution. It would need some thought of how to express the output of the said stochastic function as samples from some distribution. Try to write down what are the input and the output of the original function, and then rewrite it into something like p(output | input)

5

Were you able to rewrite your model as a custom distribution? I am facing a similar issue with a model. I was looking at trying to use an approach similar to modelling a non-homogeneous HMM to get around this issue, but with no success.

6

Hi discordiaZA, I haven’t tried writing a custom distribution yet and decided to start by making the choice deterministic, using p_high.argmax() like junpenglao suggested. It is not working yet (see here), but I’m hoping that I’ll get it to work soon. What is the issue with your model?