nsteplstm_sample.py

import numpy as np
import chainer.links as L
import chainer.functions as F
from chainer import Variable


def sequence_embed(embed, xs):
    x_len = [len(x) for x in xs]
    x_section = np.cumsum(x_len[:-1])
    ex = embed(F.concat(xs, axis=0))
    exs = F.split_axis(ex, x_section, 0)
    return exs


# 適当なニューラルネット (注意: 本来はこのような小規模なネットワークは組みません)
embed = L.EmbedID(10, 3)
lstm = L.NStepLSTM(2, 3, 3, 0.5)

# 各文の単語 ID 列を Variable でラップする
xs = [
  Variable(np.array([0, 5, 7, 1], dtype=np.int32)),
  Variable(np.array([0, 4, 3, 6, 7, 1], dtype=np.int32))
]

# 正常に処理される
emb_xs = sequence_embed(embed, xs)

hy, cy, ys = lstm(None, None, emb_xs)