keras/src/backend/tensorflow/numpy.py/expand_dims
def expand_dims(x, axis):
    x = convert_to_tensor(x)
    axis = to_tuple_or_list(axis)
    out_ndim = len(x.shape) + len(axis)
    axis = sorted([canonicalize_axis(a, out_ndim) for a in axis])
    if isinstance(x, tf.SparseTensor):
        from keras.src.ops.operation_utils import (
            compute_expand_dims_output_shape,
        )

        output_shape = compute_expand_dims_output_shape(x.shape, axis)
        for a in axis:
            x = tf.sparse.expand_dims(x, a)
        x.set_shape(output_shape)
        return x
    for a in axis:
        x = tf.expand_dims(x, a)
    return x