keras/src/backend/tensorflow/numpy.py/take
def take(x, indices, axis=None):
    if isinstance(indices, tf.SparseTensor):
        if x.dtype not in (tf.float16, tf.float32, tf.float64, tf.bfloat16):
            warnings.warn(
                "`take` with the TensorFlow backend does not support "
                f"`x.dtype={x.dtype}` when `indices` is a sparse tensor; "
                "densifying `indices`."
            )
            return take(x, convert_to_tensor(indices, sparse=False), axis=axis)
        if axis is None:
            x = tf.reshape(x, (-1,))
        elif axis != 0:
            warnings.warn(
                "`take` with the TensorFlow backend does not support "
                f"`axis={axis}` when `indices` is a sparse tensor; "
                "densifying `indices`."
            )
            return take(x, convert_to_tensor(indices, sparse=False), axis=axis)
        output = tf.nn.safe_embedding_lookup_sparse(
            embedding_weights=tf.convert_to_tensor(x),
            sparse_ids=tf.sparse.expand_dims(indices, axis=-1),
            default_id=0,
        )
        output.set_shape(indices.shape + output.shape[len(indices.shape) :])
        return output

    x = convert_to_tensor(x)
    indices = convert_to_tensor(indices)
    if axis is None:
        x = tf.reshape(x, [-1])
        axis = 0
    # Correct the indices using "fill" mode which is the same as in jax
    indices = tf.where(
        indices < 0,
        indices + tf.cast(tf.shape(x)[axis], indices.dtype),
        indices,
    )
    return tf.gather(x, indices, axis=axis)