def reshape(x, num_devices=2):
x_shape = list(x.shape)
batch_size = x_shape[0]
batch_size_per_device = batch_size // num_devices
# New shape.
new_shape_prefix = [num_devices, batch_size_per_device]
return tf_np.reshape(x, new_shape_prefix + x_shape[1:])
def scan_reference(f, init, xs):
carry = init
ys = []
for x in xs:
(carry, y) = f(carry, x)
ys.append(tf_np.reshape(y, (1,) + y.shape))
ys = tf_np.concatenate(ys, 0)
return carry, ys
def losses(scan, c, xs):
c, ys = scan(f, c, xs)
return tf_np.concatenate(tf.nest.flatten(tf.nest.map_structure(
lambda a: tf_np.reshape(a, [-1]), (c, ys))))
