def _transform_values( self, aux_values: Dict[str, torch.Tensor], ) -> Tuple[Dict[str, torch.Tensor], Union[float, torch.Tensor]]: # Learnably transform auxiliary values to user-facing values. values = {} log_densities = defaultdict(float) compute_density = am_i_wrapped() and poutine.get_mask() is not False for name, site in self._factors.items(): if site["is_observed"]: continue loc = deep_getattr(self.locs, name) scale = deep_getattr(self.scales, name) unconstrained = aux_values[name] * scale + loc # Transform to constrained space. transform = biject_to(site["fn"].support) values[name] = transform(unconstrained) if compute_density: assert transform.codomain.event_dim == site["fn"].event_dim log_densities[name] = transform.inv.log_abs_det_jacobian( values[name], unconstrained ) - scale.log().reshape(site["fn"].batch_shape + (-1,)).sum(-1) return values, log_densities
def sample(name, fn, *args, **kwargs): """ Calls the stochastic function ``fn`` with additional side-effects depending on ``name`` and the enclosing context (e.g. an inference algorithm). See `Introduction to Pyro <http://pyro.ai/examples/intro_long.html>`_ for a discussion. :param name: name of sample :param fn: distribution class or function :param obs: observed datum (optional; should only be used in context of inference) optionally specified in kwargs :param ~torch.Tensor obs_mask: Optional boolean tensor mask of shape broadcastable with ``fn.batch_shape``. If provided, events with mask=True will be conditioned on ``obs`` and remaining events will be imputed by sampling. This introduces a latent sample site named ``name + "_unobserved"`` which should be used by guides. :type obs_mask: bool or ~torch.Tensor :param dict infer: Optional dictionary of inference parameters specified in kwargs. See inference documentation for details. :returns: sample """ # Transform obs_mask into multiple sample statements. obs = kwargs.pop("obs", None) obs_mask = kwargs.pop("obs_mask", None) if obs_mask is not None: return _masked_observe(name, fn, obs, obs_mask, *args, **kwargs) # Check if stack is empty. # if stack empty, default behavior (defined here) infer = kwargs.pop("infer", {}).copy() is_observed = infer.pop("is_observed", obs is not None) if not am_i_wrapped(): if obs is not None and not infer.get("_deterministic"): warnings.warn( "trying to observe a value outside of inference at " + name, RuntimeWarning, ) return obs return fn(*args, **kwargs) # if stack not empty, apply everything in the stack? else: # initialize data structure to pass up/down the stack msg = { "type": "sample", "name": name, "fn": fn, "is_observed": is_observed, "args": args, "kwargs": kwargs, "value": obs, "infer": infer, "scale": 1.0, "mask": None, "cond_indep_stack": (), "done": False, "stop": False, "continuation": None, } # apply the stack and return its return value apply_stack(msg) return msg["value"]
def param(name, *args, **kwargs): """ Saves the variable as a parameter in the param store. To interact with the param store or write to disk, see `Parameters <parameters.html>`_. :param name: name of parameter :returns: parameter """ if not am_i_wrapped(): return _PYRO_PARAM_STORE.get_param(name, *args, **kwargs) else: msg = { "type": "param", "name": name, "args": args, "kwargs": kwargs, "infer": {}, "scale": 1.0, "cond_indep_stack": (), "value": None, "done": False, "stop": False, "continuation": None } # apply the stack and return its return value apply_stack(msg) return msg["value"]
def param(name, *args, **kwargs): """ Saves the variable as a parameter in the param store. To interact with the param store or write to disk, see `Parameters <parameters.html>`_. :param name: name of parameter :returns: parameter """ if not am_i_wrapped(): return _PYRO_PARAM_STORE.get_param(name, *args, **kwargs) else: msg = { "type": "param", "name": name, "args": args, "kwargs": kwargs, "infer": {}, "scale": 1.0, "cond_indep_stack": (), "value": None, "done": False, "stop": False, "continuation": None } # apply the stack and return its return value apply_stack(msg) return msg["value"]
def _sample_aux_values(self, *, temperature: float) -> Dict[str, torch.Tensor]: funsor = _import_funsor() # Convert torch to funsor. particle_plates = frozenset(get_plates()) plate_to_dim = self._funsor_plate_to_dim.copy() plate_to_dim.update({f.name: f.dim for f in particle_plates}) factors = {} for d, inputs in self._funsor_factor_inputs.items(): batch_shape = torch.Size( p.size for p in sorted(self._plates[d], key=lambda p: p.dim) ) white_vec = deep_getattr(self.white_vecs, d) prec_sqrt = deep_getattr(self.prec_sqrts, d) factors[d] = funsor.gaussian.Gaussian( white_vec=white_vec.reshape(batch_shape + white_vec.shape[-1:]), prec_sqrt=prec_sqrt.reshape(batch_shape + prec_sqrt.shape[-2:]), inputs=inputs, ) # Perform Gaussian tensor variable elimination. if temperature == 1: samples, log_prob = _try_possibly_intractable( funsor.recipes.forward_filter_backward_rsample, factors=factors, eliminate=self._funsor_eliminate, plates=frozenset(plate_to_dim), sample_inputs={f.name: funsor.Bint[f.size] for f in particle_plates}, ) else: samples, log_prob = _try_possibly_intractable( funsor.recipes.forward_filter_backward_precondition, factors=factors, eliminate=self._funsor_eliminate, plates=frozenset(plate_to_dim), ) # Substitute noise. sample_shape = torch.Size(f.size for f in particle_plates) noise = torch.randn(sample_shape + log_prob.inputs["aux"].shape) noise.mul_(temperature) aux = funsor.Tensor(noise)[tuple(f.name for f in particle_plates)] with funsor.interpretations.memoize(): samples = {k: v(aux=aux) for k, v in samples.items()} log_prob = log_prob(aux=aux) # Convert funsor to torch. if am_i_wrapped() and poutine.get_mask() is not False: log_prob = funsor.to_data(log_prob, name_to_dim=plate_to_dim) pyro.factor(f"_{self._pyro_name}_latent", log_prob, has_rsample=True) samples = { k: funsor.to_data(v, name_to_dim=plate_to_dim) for k, v in samples.items() } return samples
def sample(name, fn, *args, **kwargs): """ Calls the stochastic function `fn` with additional side-effects depending on `name` and the enclosing context (e.g. an inference algorithm). See `Intro I <http://pyro.ai/examples/intro_part_i.html>`_ and `Intro II <http://pyro.ai/examples/intro_part_ii.html>`_ for a discussion. :param name: name of sample :param fn: distribution class or function :param obs: observed datum (optional; should only be used in context of inference) optionally specified in kwargs :param dict infer: Optional dictionary of inference parameters specified in kwargs. See inference documentation for details. :returns: sample """ obs = kwargs.pop("obs", None) infer = kwargs.pop("infer", {}).copy() # check if stack is empty # if stack empty, default behavior (defined here) if not am_i_wrapped(): if obs is not None: warnings.warn( "trying to observe a value outside of inference at " + name, RuntimeWarning) return obs return fn(*args, **kwargs) # if stack not empty, apply everything in the stack? else: # initialize data structure to pass up/down the stack msg = { "type": "sample", "name": name, "fn": fn, "is_observed": False, "args": args, "kwargs": kwargs, "value": None, "infer": infer, "scale": 1.0, "mask": None, "cond_indep_stack": (), "done": False, "stop": False, "continuation": None } # handle observation if obs is not None: msg["value"] = obs msg["is_observed"] = True # apply the stack and return its return value apply_stack(msg) return msg["value"]
def __iter__(self): if not am_i_wrapped(): for i in self.subsample: yield i if isinstance(i, numbers.Number) else i.item() else: indep_context = poutine.indep(name=self.name, size=self.subsample_size) with poutine.scale(scale=self.size / self.subsample_size): for i in self.subsample: indep_context.next_context() with indep_context: # convert to python numeric type as functions like torch.ones(*args) # do not work with dim 0 torch.Tensor instances. yield i if isinstance(i, numbers.Number) else i.item()
def __enter__(self): self._wrapped = am_i_wrapped() self.dim = _DIM_ALLOCATOR.allocate(self.name, self.dim) if self._wrapped: try: self._scale_messenger = poutine.scale(scale=self.size / self.subsample_size) self._indep_messenger = poutine.indep(name=self.name, size=self.subsample_size, dim=self.dim) self._scale_messenger.__enter__() self._indep_messenger.__enter__() except BaseException: _DIM_ALLOCATOR.free(self.name, self.dim) raise return self.subsample
def __iter__(self): if not am_i_wrapped(): for i in self.subsample: yield i if isinstance(i, numbers.Number) else i.item() else: indep_context = poutine.indep(name=self.name, size=self.subsample_size) with poutine.scale(scale=self.size / self.subsample_size): for i in self.subsample: indep_context.next_context() with indep_context: # convert to python numeric type as functions like torch.ones(*args) # do not work with dim 0 torch.Tensor instances. yield i if isinstance(i, numbers.Number) else i.item()
def sample(name, fn, *args, **kwargs): """ Calls the stochastic function `fn` with additional side-effects depending on `name` and the enclosing context (e.g. an inference algorithm). See `Intro I <http://pyro.ai/examples/intro_part_i.html>`_ and `Intro II <http://pyro.ai/examples/intro_part_ii.html>`_ for a discussion. :param name: name of sample :param fn: distribution class or function :param obs: observed datum (optional; should only be used in context of inference) optionally specified in kwargs :param dict infer: Optional dictionary of inference parameters specified in kwargs. See inference documentation for details. :returns: sample """ obs = kwargs.pop("obs", None) infer = kwargs.pop("infer", {}) # check if stack is empty # if stack empty, default behavior (defined here) if not am_i_wrapped(): if obs is not None: warnings.warn("trying to observe a value outside of inference at " + name, RuntimeWarning) return obs return fn(*args, **kwargs) # if stack not empty, apply everything in the stack? else: # initialize data structure to pass up/down the stack msg = { "type": "sample", "name": name, "fn": fn, "is_observed": False, "args": args, "kwargs": kwargs, "value": None, "infer": infer, "scale": 1.0, "cond_indep_stack": (), "done": False, "stop": False, "continuation": None } # handle observation if obs is not None: msg["value"] = obs msg["is_observed"] = True # apply the stack and return its return value apply_stack(msg) return msg["value"]
def __enter__(self): self._wrapped = am_i_wrapped() self.dim = _DIM_ALLOCATOR.allocate(self.name, self.dim) if self._wrapped: try: self._scale_messenger = poutine.scale(scale=self.size / self.subsample_size) self._indep_messenger = poutine.indep(name=self.name, size=self.subsample_size, dim=self.dim) self._scale_messenger.__enter__() self._indep_messenger.__enter__() except BaseException: _DIM_ALLOCATOR.free(self.name, self.dim) raise return self.subsample
def forward(self, zs): embedding = self.pre_recurrence_linear(zs) hiddens = [None, None] teacher = None if runtime.am_i_wrapped() and\ isinstance(runtime._PYRO_STACK[-1], ConditionMessenger): data = runtime._PYRO_STACK[-1].data if '$%s$' % self._smiles_name in data: teacher = data['$%s$' % self._smiles_name] probs = [] for i in range(self._max_len): hiddens[0] = self.recurrence1(embedding, hiddens[0]) hiddens[1] = self.recurrence2(hiddens[0], hiddens[1]) embedding = self.decoder(hiddens[1]) probs.append(embedding) if teacher is not None: embedding = teacher[:, i] probs = torch.stack(probs, dim=1) logits_categorical = dist.OneHotCategorical(probs=probs).to_event(1) return pyro.sample('$%s$' % self._smiles_name, logits_categorical)