def _value_at(x, idx):
    """
    input x: 1D tensor (vector).
    return value at index idx. x[idx].
    """
    assert x.rank == 1
    return mb.slice_by_index(x=x, begin=[idx], end=[0], squeeze_mask=[True])

    input x: 1D tensor (vector).
    return value at index idx. x[idx].
    
coremltools/converters/mil/frontend/tensorflow/ops.py