def _setup_prototype(self, *args, **kwargs):
super()._setup_prototype(*args, **kwargs)
for name, site in self.prototype_trace.items():
if site["type"] != "sample" or site["is_observed"]:
continue
event_dim = site["fn"].event_dim + jnp.ndim(self._init_locs[name]) - jnp.ndim(site["value"])
self._event_dims[name] = event_dim
# If subsampling, repeat init_value to full size.
for frame in site["cond_indep_stack"]:
full_size = self._prototype_frame_full_sizes[frame.name]
if full_size != frame.size:
dim = frame.dim - event_dim
self._init_locs[name] = periodic_repeat(self._init_locs[name], full_size, dim)
def _setup_prototype(self, *args, **kwargs):
super()._setup_prototype(*args, **kwargs)
with numpyro.handlers.block():
self._init_locs = {
k: v
for k, v in self._postprocess_fn(self._init_locs).items()
if k in self._init_locs
}
for name, site in self.prototype_trace.items():
if site["type"] != "sample" or site["is_observed"]:
continue
event_dim = site["fn"].event_dim
self._event_dims[name] = event_dim
# If subsampling, repeat init_value to full size.
for frame in site["cond_indep_stack"]:
full_size = self._prototype_frame_full_sizes[frame.name]
if full_size != frame.size:
dim = frame.dim - event_dim
self._init_locs[name] = periodic_repeat(
self._init_locs[name], full_size, dim)