permutes the dimensions of the input according to a given pattern useful e g connecting rnns and convnets args dims tuple of integers permutation pattern does not include the batch dimension indexing starts at 1 for instance 2 1 permutes the first and second dimensions of the input input shape arbitrary output shape same as the input shape but with the dimensions reordered according to the specified pattern example x keras inputshape10 64 y keras layers permute2 1x y shape none 64 10 permutes the dimensions of the input according to a given pattern useful e g connecting rnns and convnets args dims tuple of integers permutation pattern does not include the batch dimension indexing starts at 1 for instance 2 1 permutes the first and second dimensions of the input input shape arbitrary output shape same as the input shape but with the dimensions re ordered according to the specified pattern example x keras input shape 10 64 y keras layers permute 2 1 x y shape none 64 10
from keras import ops
from keras.api_export import keras_export
from keras.backend.common.keras_tensor import KerasTensor
from keras.layers.input_spec import InputSpec
from keras.layers.layer import Layer


@keras_export("keras.layers.Permute")
class Permute(Layer):
    

    def __init__(self, dims, **kwargs):
        super().__init__(**kwargs)
        self.dims = tuple(dims)
        if sorted(dims) != list(range(1, len(dims) + 1)):
            raise ValueError(
                "Invalid permutation argument `dims` for Permute Layer. "
                "The set of indices in `dims` must be consecutive and start "
                f"from 1. Received dims={dims}"
            )
        self.input_spec = InputSpec(ndim=len(self.dims) + 1)

    def compute_output_shape(self, input_shape):
        output_shape = [input_shape[0]]
        for dim in self.dims:
            output_shape.append(input_shape[dim])
        return tuple(output_shape)

    def compute_output_spec(self, inputs):
        output_shape = self.compute_output_shape(inputs.shape)
        return KerasTensor(
            shape=output_shape, dtype=inputs.dtype, sparse=inputs.sparse
        )

    def call(self, inputs):
        return ops.transpose(inputs, axes=(0,) + self.dims)

    def get_config(self):
        config = {"dims": self.dims}
        base_config = super().get_config()
        return {**base_config, **config}