use .mask() in air example
@@ -14,7 +14,6 @@ import torch.nn.functional as F
import pyro
import pyro.distributions as dist
-import pyro.poutine as poutine
from modules import Identity, Encoder, Decoder, MLP, Predict
@@ -133,19 +132,20 @@ class AIR(nn.Module):
# will be added to its output image. We can't
# straight-forwardly avoid generating further objects, so
# instead we zero out the log_prob_sum of future choices.
- sample_mask = z_pres.squeeze(-1) if self.use_masking else 1
- with poutine.scale(None, sample_mask):
+ sample_mask = z_pres if self.use_masking else torch.tensor(1.0)
# Sample attention window position.
z_where = pyro.sample('z_where_{}'.format(t),
dist.Normal(self.z_where_loc_prior.expand(n, self.z_where_size),
self.z_where_scale_prior.expand(n, self.z_where_size))
+ .mask(sample_mask)
.reshape(extra_event_dims=1))
# Sample latent code for contents of the attention window.
z_what = pyro.sample('z_what_{}'.format(t),
dist.Normal(self.prototype.new_zeros([n, self.z_what_size]),
self.prototype.new_ones([n, self.z_what_size]))
+ .mask(sample_mask)
.reshape(extra_event_dims=1))
# Map latent code to pixel space.
@@ -235,10 +235,12 @@ class AIR(nn.Module):
dist.Bernoulli(z_pres_p * prev.z_pres).reshape(extra_event_dims=1),
infer=dict(baseline=dict(baseline_value=bl_value.squeeze(-1))))
- with poutine.scale(None, z_pres.squeeze(-1) if self.use_masking else 1):
+ sample_mask = z_pres if self.use_masking else torch.tensor(1.0)
+
z_where = pyro.sample('z_where_{}'.format(t),
dist.Normal(z_where_loc + self.z_where_loc_prior,
z_where_scale * self.z_where_scale_prior)
+ .mask(sample_mask)
.reshape(extra_event_dims=1))
# Figure 2 of [1] shows x_att depending on z_where and h,
@@ -251,6 +253,7 @@ class AIR(nn.Module):
z_what = pyro.sample('z_what_{}'.format(t),
dist.Normal(z_what_loc, z_what_scale)
+ .mask(sample_mask)
.reshape(extra_event_dims=1))
return GuideState(h=h, c=c, bl_h=bl_h, bl_c=bl_c, z_pres=z_pres, z_where=z_where, z_what=z_what)
