Dear community,
In trying to use VI to speed up a model that I have previously made work (first code block below) in the MCMC framework (took several days to run and the results were good), I run into this error of “ValueError: Multiple update expressions found for the variable RandomGeneratorSharedVariable”, which I suspect to have something to do with the way I set up minibatches (second code block below).
I have set up the total_size according to this thread https://discourse.pymc.io/t/possible-extensions-of-total-size/36.
When it comes to setting up pm.minibatch on multi-dimensional variables, it was inconclusive from fearrine’s commends in this thread https://discourse.pymc.io/t/how-to-make-minibatch-for-multi-dimensional-data/5033/4. I couldn’t find any other examples or discussion on this matter.
Details on the model and variables: the observed samples are non-negative integers (hence the ZIP likelihood) under 2 dimensions, one of which (dimA) is in size of hundreds of thousands (target of minibatch) and the other (dimB) is 40. Also 3 categorical covariates are involved (c1 has 12 possible values, c2 has 2 and c3 has 2). 40 samples in dimB share the same set of covariates, while samples in dimA are latently grouped within a finite number of sets of covariates (12 X 2 X 2=48 combos) but also allow to vary around these finite (48) centers with variations also modeled.
MCMC code that works:
with pm.Model() as m:
m.add_coord('nC1', dataDict['dataDF']['nC1'].unique(), mutable=True)
m.add_coord('nC2', dataDict['dataDF']['nC2'].unique(), mutable=True)
m.add_coord('nC3', dataDict['dataDF']['nC3'].unique(), mutable=True)
m.add_coord('dimA', ..., mutable=True)
m.add_coord('dimB', ...), mutable=True)
TD_obs = pm.MutableData('TD_obs', dataDict['TD'], dims=('dimB', 'dimA'))
AD_obs = pm.MutableData('AD_obs', dataDict['AD'], dims=('dimB', 'dimA'))
c1_idx = pm.MutableData('c1_idx', dataDict['dataDF']['context'].values, dims='dimA')
c2_idx = pm.MutableData('c2_idx', dataDict['dataDF']['segDup'].values, dims='dimA')
c3_idx = pm.MutableData('c3_idx', dataDict['dataDF']['mapa'].values, dims='dimA')
mu_bc = pm.TruncatedNormal('mu_bc', mu=8, sigma=3, lower=4, shape=1)
std_bc = pm.HalfNormal('std_bc', sigma=1, shape=1)
mu_c1 = pm.TruncatedNormal('mu_c1', mu=mu_bc, sigma=std_bc, lower=4, dims='nC1')
mu_c2 = pm.Normal('mu_c2', mu=0, sigma=1, dims='nC2')
mu_c3 = pm.Normal('mu_c3', mu=0, sigma=1, dims='nC3')
mu_dimA = pm.Deterministic('mu_dimA', mu_c1[c1_idx] + mu_c2[c2_idx] + mu_c3[c3_idx], dims='dimA')
std_shared = pm.HalfNormal('std_dimA', sigma=2, shape=1)
ER_dimA = pm.Gamma('ER_dimA', alpha=mu_dimA ** 2 / std_shared ** 2, beta=mu_dimA / std_shared ** 2, dims='dimA')
psi_dimA = pm.Beta('psi_dimA', alpha=2, beta=5, dims='dimA')
AD_predicted = pm.ZeroInflatedBinomial('AD_predicted', psi=psi_dimA, n=TD_obs,
p=pm.invlogit(-ER_dimA), observed=AD_obs)
VI with minibatches:
with pm.Model() as m:
posBatchSize = 100
lenDimB = 40
m.add_coord('nC1', dataDict['dataDF']['nC1'].unique(), mutable=True)
m.add_coord('nC2', dataDict['dataDF']['nC2'].unique(), mutable=True)
m.add_coord('nC3', dataDict['dataDF']['nC3'].unique(), mutable=True)
# these two observed variables have 2 dimensions (lenDimB, lenPos)
TD_obs_m = pm.Minibatch(dataDict['TD'], batch_size=[(lenDimB, posBatchSize)])
AD_obs_m = pm.Minibatch(dataDict['AD'], batch_size=[(lenDimB, posBatchSize)])
# these three covariates have 1 dimension (lenPos,)
c1_idx_m = pm.Minibatch(dataDict['dataDF']['c1'].values, batch_size=posBatchSize)
c2_idx_m = pm.Minibatch(dataDict['dataDF']['c2'].values, batch_size=posBatchSize)
c3_idx_m = pm.Minibatch(dataDict['dataDF']['c3'].values, batch_size=posBatchSize) (nPos,)
mu_bc = pm.TruncatedNormal('mu_bc', mu=8, sigma=3, lower=4, shape=1)
std_bc = pm.HalfNormal('std_bc', sigma=1, shape=1)
mu_c1 = pm.TruncatedNormal('mu_c1', mu=mu_bc, sigma=std_bc, lower=4, dims='nC1')
mu_c2 = pm.Normal('mu_c2', mu=0, sigma=1, dims='nC2')
mu_c3 = pm.Normal('mu_c3', mu=0, sigma=1, dims='nC3')
mu_dimA = pm.Deterministic('mu_dimA', mu_c1[c1_idx] + mu_c2[c2_idx] + mu_c3[c3_idx], dims='dimA')
std_shared = pm.HalfNormal('std_dimA', sigma=2, shape=1)
ER_dimA = pm.Gamma('ER_dimA', alpha=mu_dimA ** 2 / std_shared ** 2, beta=mu_dimA / std_shared ** 2, dims='dimA')
psi_dimA = pm.Beta('psi_dimA', alpha=2, beta=5, dims='dimA')
AD_predicted = pm.ZeroInflatedBinomial(
'AD_predicted', psi=psi_dimA, n=TD_obs_m, p=pm.invlogit(-ER_dimA), observed=AD_obs_m, total_size=dataDict['AD'].shape)
approx = pm.fit(1_000, callbacks=[pm.callbacks.CheckParametersConvergence(tolerance=1e-4)])
it gave an error:
/root/anaconda3/envs/pm5/lib/python3.11/site-packages/pymc/logprob/joint_logprob.py:167: UserWarning: Found a random variable that was neither among the observations nor the conditioned variables: [integers_rv{0, (0, 0), int64, False}.0, integers_rv{0, (0, 0), int64, False}.out]
warnings.warn(
/root/anaconda3/envs/pm5/lib/python3.11/site-packages/pymc/logprob/joint_logprob.py:167: UserWarning: Found a random variable that was neither among the observations nor the conditioned variables: [integers_rv{0, (0, 0), int64, False}.0, integers_rv{0, (0, 0), int64, False}.out]
warnings.warn(
/root/anaconda3/envs/pm5/lib/python3.11/site-packages/pymc/logprob/joint_logprob.py:167: UserWarning: Found a random variable that was neither among the observations nor the conditioned variables: [integers_rv{0, (0, 0), int64, False}.0, integers_rv{0, (0, 0), int64, False}.out]
warnings.warn(
/root/anaconda3/envs/pm5/lib/python3.11/site-packages/pymc/logprob/joint_logprob.py:167: UserWarning: Found a random variable that was neither among the observations nor the conditioned variables: [integers_rv{0, (0, 0), int64, False}.0, integers_rv{0, (0, 0), int64, False}.out]
warnings.warn(
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Cell In[87], line 2
1 with m:
----> 2 approx = pm.fit(1_000, callbacks=[pm.callbacks.CheckParametersConvergence(tolerance=1e-4)])
4 # idata_advi = approx.sample(500)
File ~/anaconda3/envs/pm5/lib/python3.11/site-packages/pymc/variational/inference.py:747, in fit(n, method, model, random_seed, start, start_sigma, inf_kwargs, **kwargs)
745 else:
746 raise TypeError(f"method should be one of {set(_select.keys())} or Inference instance")
--> 747 return inference.fit(n, **kwargs)
File ~/anaconda3/envs/pm5/lib/python3.11/site-packages/pymc/variational/inference.py:138, in Inference.fit(self, n, score, callbacks, progressbar, **kwargs)
136 callbacks = []
137 score = self._maybe_score(score)
--> 138 step_func = self.objective.step_function(score=score, **kwargs)
139 if progressbar:
140 progress = progress_bar(range(n), display=progressbar)
File ~/anaconda3/envs/pm5/lib/python3.11/site-packages/pytensor/configparser.py:47, in _ChangeFlagsDecorator.__call__.<locals>.res(*args, **kwargs)
44 @wraps(f)
45 def res(*args, **kwargs):
46 with self:
---> 47 return f(*args, **kwargs)
File ~/anaconda3/envs/pm5/lib/python3.11/site-packages/pymc/variational/opvi.py:387, in ObjectiveFunction.step_function(self, obj_n_mc, tf_n_mc, obj_optimizer, test_optimizer, more_obj_params, more_tf_params, more_updates, more_replacements, total_grad_norm_constraint, score, fn_kwargs)
385 seed = self.approx.rng.randint(2**30, dtype=np.int64)
386 if score:
--> 387 step_fn = compile_pymc([], updates.loss, updates=updates, random_seed=seed, **fn_kwargs)
388 else:
389 step_fn = compile_pymc([], [], updates=updates, random_seed=seed, **fn_kwargs)
File ~/anaconda3/envs/pm5/lib/python3.11/site-packages/pymc/pytensorf.py:1083, in compile_pymc(inputs, outputs, random_seed, mode, **kwargs)
1047 """Use ``pytensor.function`` with specialized pymc rewrites always enabled.
1048
1049 This function also ensures shared RandomState/Generator used by RandomVariables
(...)
1079 is set to False.
1080 """
1081 # Create an update mapping of RandomVariable's RNG so that it is automatically
1082 # updated after every function call
-> 1083 rng_updates = collect_default_updates(inputs, outputs)
1085 # We always reseed random variables as this provides RNGs with no chances of collision
1086 if rng_updates:
File ~/anaconda3/envs/pm5/lib/python3.11/site-packages/pymc/pytensorf.py:1036, in collect_default_updates(inputs, outputs)
1032 # When a variable has multiple outputs, it will be called twice with the same
1033 # update expression. We don't want to raise in that case, only if the update
1034 # expression in different from the one already registered
1035 elif rng_updates[rng] is not update:
-> 1036 raise ValueError(f"Multiple update expressions found for the variable {rng}")
1037 return rng_updates
ValueError: Multiple update expressions found for the variable RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F0B95FC5000>)
I can provide further information if necessary. Thanks in advance for any advices.