def take_along_axis(x, indices, axis=None):
    """Select values from `x` at the 1-D `indices` along the given axis.

    Args:
        x: Source tensor.
        indices: The indices of the values to extract.
        axis: The axis over which to select values. By default, the flattened
            input tensor is used.

    Returns:
        The corresponding tensor of values.
    """
    if any_symbolic_tensors((x, indices)):
        return TakeAlongAxis(axis=axis).symbolic_call(x, indices)
    return backend.numpy.take_along_axis(x, indices, axis=axis)
Select values from `x` at the 1-D `indices` along the given axis.

    Args:
        x: Source tensor.
        indices: The indices of the values to extract.
        axis: The axis over which to select values. By default, the flattened
            input tensor is used.

    Returns:
        The corresponding tensor of values.
    
keras/src/ops/numpy.py