def __call__(self, sequences):
"""Embeds int or string sequences."""
if isinstance(sequences[0], str):
sequences = _encode_string_sequences(sequences,
domain=self._domain,
length=self._length)
return utils.batch_apply(self._embed_fn, sequences, self._batch_size)
def test_batch_apply(self, batch_size, num_inputs):
def fn(inputs):
return np.power(inputs + 1, 2)
def batch_fn(batched_inputs):
if len(batched_inputs) != batch_size:
raise ValueError('fn() called with a batch that is '
'the wrong size (%d vs. %d).' % (len(batched_inputs),
batch_size))
return fn(batched_inputs)
inputs = np.stack([np.arange(num_inputs), -np.arange(num_inputs)], axis=1)
unbatched_output = fn(inputs)
batched_output = utils.batch_apply(batch_fn, inputs, batch_size)
np.testing.assert_array_equal(unbatched_output, batched_output)