def rearrange(tensor, pattern, **axes_lengths):
    """Rearranges the axes of a Keras tensor according to a specified pattern,
    einops-style.

    Args:
        tensor: Input Keras tensor.
        pattern: String describing the rearrangement in einops notation.
        **axes_lengths: Keyword arguments specifying lengths of axes
            when axes decomposition is used.

    Returns:
        Tensor: A Keras tensor with rearranged axes.

    Follows the logic of:

    1. If decomposition is needed, reshape to match decomposed dimensions.
    2. Permute known and inferred axes to match the form of the output.
    3. Reshape to match the desired output shape.


    Example Usage:

    ```
    >>> import numpy as np
    >>> from keras.ops import rearrange
    >>> images = np.random.rand(32, 30, 40, 3) # BHWC format

    # Reordering to BCHW
    >>> rearrange(images, 'b h w c -> b c h w').shape
    TensorShape([32, 3, 30, 40])

    # "Merge" along first axis - concat images from a batch
    >>> rearrange(images, 'b h w c -> (b h) w c').shape
    TensorShape([960, 40, 3])

    # "Merge" along second axis - concat images horizontally
    >>> rearrange(images, 'b h w c -> h (b w) c').shape
    TensorShape([30, 1280, 3])

    # Flatten images into a CHW vector
    >>> rearrange(images, 'b h w c -> b (c h w)').shape
    TensorShape([32, 3600])

    # Decompose H and W axes into 4 smaller patches
    >>> rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', h1=2, w1=2).shape
    TensorShape([128, 15, 20, 3])

    # Space-to-depth decomposition of input axes
    >>> rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, w1=2).shape
    TensorShape([32, 15, 20, 12])
    ```
    """  # noqa: E501

    if any_symbolic_tensors((tensor,)):
        return Rearrange().symbolic_call(tensor, pattern, **axes_lengths)

    # Split the input and output patterns
    input_pattern, output_pattern = re.split(r"\s*->\s*", pattern)
    input_axes = re.findall(r"\w+|\(.*?\)", input_pattern)
    output_axes = re.findall(r"\w+|\(.*?\)", output_pattern)
    input_shape = shape(tensor)

    # Create axes map, and flattened output group
    axes_map = _create_axes_map(input_axes, input_shape, axes_lengths)
    grouped_output_axes = _create_grouped_axes(output_axes)
    flattened_output_axes = _flatten_group(grouped_output_axes)

    # 1. Axes decomposition
    decomposed_shapes = _compute_decomposed_shape(
        input_axes, axes_lengths, axes_map
    )
    if decomposed_shapes != tensor.shape:
        tensor = reshape(tensor, decomposed_shapes)

    # 2. Transpose to match target shape
    permute_order = _get_transpose_order(input_axes, flattened_output_axes)
    tensor = transpose(tensor, permute_order)

    # 3. Reshape to final target shape
    output_shape = _compute_output_shape(axes_map, grouped_output_axes)
    tensor = reshape(tensor, output_shape)

    return tensor
Rearranges the axes of a Keras tensor according to a specified pattern,
    einops-style.

    Args:
        tensor: Input Keras tensor.
        pattern: String describing the rearrangement in einops notation.
        **axes_lengths: Keyword arguments specifying lengths of axes
            when axes decomposition is used.

    Returns:
        Tensor: A Keras tensor with rearranged axes.

    Follows the logic of:

    1. If decomposition is needed, reshape to match decomposed dimensions.
    2. Permute known and inferred axes to match the form of the output.
    3. Reshape to match the desired output shape.


    Example Usage:

    ```
    >>> import numpy as np
    >>> from keras.ops import rearrange
    >>> images = np.random.rand(32, 30, 40, 3) # BHWC format

    # Reordering to BCHW
    >>> rearrange(images, 'b h w c -> b c h w').shape
    TensorShape([32, 3, 30, 40])

    # "Merge" along first axis - concat images from a batch
    >>> rearrange(images, 'b h w c -> (b h) w c').shape
    TensorShape([960, 40, 3])

    # "Merge" along second axis - concat images horizontally
    >>> rearrange(images, 'b h w c -> h (b w) c').shape
    TensorShape([30, 1280, 3])

    # Flatten images into a CHW vector
    >>> rearrange(images, 'b h w c -> b (c h w)').shape
    TensorShape([32, 3600])

    # Decompose H and W axes into 4 smaller patches
    >>> rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', h1=2, w1=2).shape
    TensorShape([128, 15, 20, 3])

    # Space-to-depth decomposition of input axes
    >>> rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, w1=2).shape
    TensorShape([32, 15, 20, 12])
    ```
    
keras/src/ops/einops.py