def split(self, batch_size, data):
    'TODO: Annotation\n        params:\n            batch_size: int\n            data: [B, x]\n        '
    if self.memory_net.use_rnn:
        data = tf.reshape(data, [batch_size, (- 1), tf.shape(data)[(- 1)]])
        (d, d_) = (data[:, :(- 1)], data[:, 1:])
        (d, d_) = (tf.reshape(d, [(- 1), tf.shape(d)[(- 1)]]), tf.reshape(d_, [(- 1), tf.shape(d_)[(- 1)]]))
        return (d, d_)
    else:
        return tf.split(data, num_or_size_splits=2, axis=0)
TODO: Annotation
params:
    batch_size: int
    data: [B, x]
def split(self, batch_size, data):
    'TODO: Annotation\n        params:\n            batch_size: int\n            data: [B, x]\n        '
    if self.memory_net.use_rnn:
        data = tf.reshape(data, [batch_size, (- 1), tf.shape(data)[(- 1)]])
        (d, d_) = (data[:, :(- 1)], data[:, 1:])
        (d, d_) = (tf.reshape(d, [(- 1), tf.shape(d)[(- 1)]]), tf.reshape(d_, [(- 1), tf.shape(d_)[(- 1)]]))
        return (d, d_)
    else:
        return tf.split(data, num_or_size_splits=2, axis=0)