File tree Expand file tree Collapse file tree 1 file changed +38
-0
lines changed
pymc_experimental/model/transforms Expand file tree Collapse file tree 1 file changed +38
-0
lines changed Original file line number Diff line number Diff line change @@ -246,6 +246,44 @@ def _(
246
246
return vip_rep
247
247
248
248
249
+ @_vip_reparam_node .register
250
+ def _ (
251
+ op : pm .Exponential ,
252
+ node : Apply ,
253
+ name : str ,
254
+ dims : List [Variable ],
255
+ transform : Optional [Transform ],
256
+ lam : pt .TensorVariable ,
257
+ ) -> ModelDeterministic :
258
+ rng , size , scale = node .inputs
259
+ scale_centered = scale ** lam
260
+ scale_noncentered = scale ** (1 - lam )
261
+ vip_rv_ = pm .Exponential .dist (
262
+ scale = scale_centered ,
263
+ size = size ,
264
+ rng = rng ,
265
+ )
266
+ vip_rv_value_ = vip_rv_ .clone ()
267
+ vip_rv_ .name = f"{ name } ::tau_"
268
+ if transform is not None :
269
+ vip_rv_value_ .name = f"{ vip_rv_ .name } _{ transform .name } __"
270
+ else :
271
+ vip_rv_value_ .name = vip_rv_ .name
272
+ vip_rv = model_free_rv (
273
+ vip_rv_ ,
274
+ vip_rv_value_ ,
275
+ transform ,
276
+ * dims ,
277
+ )
278
+
279
+ vip_rep_ = scale_noncentered * vip_rv
280
+
281
+ vip_rep_ .name = name
282
+
283
+ vip_rep = model_deterministic (vip_rep_ , * dims )
284
+ return vip_rep
285
+
286
+
249
287
def vip_reparametrize (
250
288
model : pm .Model ,
251
289
var_names : Sequence [str ],
You can’t perform that action at this time.
0 commit comments