def _pack_hidden(self, hidden: to.Tensor, batch_size: int = None): """ Pack the hidden state returned by the network into an 1-dim state vector. This should be the reverse operation of `_unpack_hidden`. The default implementation is defined by `default_pack_hidden`. :param hidden: hidden state as returned by the network :param batch_size: if not `None`, the result should be 2-dim and the first dim represents parts of a data batch :return: packed hidden state """ return default_pack_hidden(hidden, self.num_recurrent_layers, self._hidden_size, batch_size)
def test_hidden_state_packing_nobatch(): num_layers = 2 hidden_size = 2 batch_size = None unpacked = to.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]) # l1 # l2 packed = to.tensor([1.0, 2.0, 3.0, 4.0]) # Test unpack pu = default_unpack_hidden(packed, num_layers, hidden_size, batch_size) to.testing.assert_allclose(pu, unpacked) # Test pack up = default_pack_hidden(unpacked, num_layers, hidden_size, batch_size) to.testing.assert_allclose(up, packed)
def test_hidden_state_packing_batch(): num_layers = 2 hidden_size = 2 batch_size = 2 unpacked = to.tensor([[[1.0, 2.0], # l1, b1 [5.0, 6.0]], # l1, b2 [[3.0, 4.0], # l2, b1 [7.0, 8.0]]]) # l2, b2 packed = to.tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]) # Test unpack pu = default_unpack_hidden(packed, num_layers, hidden_size, batch_size) to.testing.assert_allclose(pu, unpacked) # Test pack up = default_pack_hidden(unpacked, num_layers, hidden_size, batch_size) to.testing.assert_allclose(up, packed)