Equivalent to `tensor[index, ...]`.

    This is a workaround for XLA requiring constant tensor indices. It works
    by producing a node representing hardcoded instructions like the following:

       if index == 0: return tensor[0]
       if index == 1: return tensor[1]
       if index == 2: return tensor[2]
       .
       .
       if index == n-1: return tensor[n-1]

    This is acceptable as long as n*size(tensor) is negligible compared to
    the rest of the computation.

def _deref(tensor: tf.Tensor, index: tf.Tensor) -> tf.Tensor:
    """Equivalent to `tensor[index, ...]`.

    This is a workaround for XLA requiring constant tensor indices. It works
    by producing a node representing hardcoded instructions like the following:

       if index == 0: return tensor[0]
       if index == 1: return tensor[1]
       if index == 2: return tensor[2]
       .
       .
       if index == n-1: return tensor[n-1]

    This is acceptable as long as n*size(tensor) is negligible compared to
    the rest of the computation.
    """
    assert tensor.shape[0] > 0
    return _deref_helper(lambda i: tensor[i, ...],
                         index,
                         0,
                         tensor.shape[0] - 1)