import tensorflow as tf


def _tf_multi_index(t, indices):
    """Partial TensorFlow implementation of Theano t[indices]."""
    # Note: this is far from a full implementation of Theano fancy
    # indexing, use with care.
    assert K._BACKEND == 'tensorflow'
    from collections import Sequence
    import tensorflow as tf

    if not isinstance(indices, Sequence):
        raise ValueError(indices)

    if len(indices) == 1:
        return tf.gather(t, indices[0])    # gather() suffices for 1d
    if K.ndim(t) == len(indices):
        # Index n-dimensional tensor with n indices: pack the indices
        # from e.g. [[i_0, i_1, ...] [j_0, j_1, ...]] to [[i_0, j_0],
        # [i_1, j_1], ...] and use gather_nd()
        # (https://www.tensorflow.org/api_docs/python/array_ops.html#gather_nd)
        # TODO: check that all i in indices have ndim n-1 
        # TODO: support broadcasting for numpy arrays with np.broadcast_to()
        #indices = tf.pack(list(indices), axis=len(indices)-1)
        indices = tf.pack(list(indices), axis=-1)
        # indices = tf.Print(indices, [indices], 'indices', summarize=100)
        return tf.gather_nd(t, indices)
    else:
