keras/src/backend/tensorflow/numpy.py/take_along_axis
def take_along_axis(x, indices, axis=None):
    from keras.src.ops.operation_utils import (
        compute_take_along_axis_output_shape,
    )

    x = convert_to_tensor(x)
    indices = convert_to_tensor(indices, "int64")
    if axis is None:
        if indices.ndim != 1:
            raise ValueError(
                "`indices` must be 1D if axis=None. "
                f"Received: indices.shape={indices.shape}"
            )
        return take_along_axis(tf.reshape(x, [-1]), indices, 0)

    # Compute the static output shape as later on, all shapes manipulations
    # use dynamic shapes.
    static_output_shape = compute_take_along_axis_output_shape(
        x.shape, indices.shape, axis
    )
    rank = x.ndim
    static_axis = axis
    axis = axis + rank if axis < 0 else axis

    # Broadcast shapes to match, ensure that the axis of interest is not
    # broadcast.
    x_shape_original = tf.shape(x, out_type=indices.dtype)
    indices_shape_original = tf.shape(indices, out_type=indices.dtype)
    x_shape = tf.tensor_scatter_nd_update(x_shape_original, [[axis]], [1])
    indices_shape = tf.tensor_scatter_nd_update(
        indices_shape_original, [[axis]], [1]
    )
    broadcasted_shape = tf.broadcast_dynamic_shape(x_shape, indices_shape)
    x_shape = tf.tensor_scatter_nd_update(
        broadcasted_shape, [[axis]], [x_shape_original[axis]]
    )
    indices_shape = tf.tensor_scatter_nd_update(
        broadcasted_shape, [[axis]], [indices_shape_original[axis]]
    )
    x = tf.broadcast_to(x, x_shape)
    indices = tf.broadcast_to(indices, indices_shape)

    # Correct the indices using "fill" mode which is the same as in jax
    indices = tf.where(indices < 0, indices + x_shape[static_axis], indices)

    x = swapaxes(x, static_axis, -1)
    indices = swapaxes(indices, static_axis, -1)

    x_shape = tf.shape(x)
    x = tf.reshape(x, [-1, x_shape[-1]])
    indices_shape = tf.shape(indices)
    indices = tf.reshape(indices, [-1, indices_shape[-1]])

    result = tf.gather(x, indices, batch_dims=1)
    result = tf.reshape(result, indices_shape)
    result = swapaxes(result, static_axis, -1)
    result.set_shape(static_output_shape)
    return result