def compute_expand_dims_output_shape(input_shape, axis):
    """Compute the output shape for the `expand_dims` operation.

    Args:
        input_shape: Input shape.
        axis: int or sequence of ints for the axis to expand.

    Returns:
        Tuple of ints: The output shape after the `expand_dims` operation.
    """
    input_shape = list(input_shape)
    if axis is None:
        axis = len(input_shape)
    axis = to_tuple_or_list(axis)
    out_ndim = len(axis) + len(input_shape)
    axis = [canonicalize_axis(a, out_ndim) for a in axis]
    shape_iter = iter(input_shape)
    new_shape = [
        1 if ax in axis else next(shape_iter) for ax in range(out_ndim)
    ]
    return tuple(new_shape)
Compute the output shape for the `expand_dims` operation.

    Args:
        input_shape: Input shape.
        axis: int or sequence of ints for the axis to expand.

    Returns:
        Tuple of ints: The output shape after the `expand_dims` operation.
    
keras/src/ops/operation_utils.py