docs: 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)