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)