def initial_temporal_state(self, *args, **kwargs):
if not hasattr(self, '_initial_temporal_state'):
state = self._propagate._ssm._cell._temporal_cell.initial_state(*args, **kwargs)
self._initial_temporal_state = nested.tile_along_newaxis(state, self._n_steps, 1)
return self._initial_temporal_state
def initial_prior_state(self, batch_size):
if not hasattr(self, '_initial_prior_state'):
prior_init_state = self._propagate.prior_init_state(batch_size, trainable=True)
self._initial_prior_state = nested.tile_along_newaxis(prior_init_state, self._n_steps, 1)
return self._initial_prior_state