def forward(
        self, channels: torch.Tensor, position: torch.Tensor, direction: torch.Tensor
    ) -> torch.Tensor:
        """
        :param channels: [batch_shape, inner_batch_shape, n_channels, height, width]
        :param position: [batch_shape, inner_batch_shape, 3, height, width]
        :param direction: [batch_shape, inner_batch_shape, 3, height, width]
        :return: [*batch_shape, out_features, height, width]
        """

        if self.posemb_version is not None:
            channels = channels.permute(0, 1, 3, 4, 2)
            position = position.permute(0, 1, 3, 4, 2)
            direction = direction.permute(0, 1, 3, 4, 2)
            channels = encode_channels(self.posemb_version, channels=channels).permute(
                0, 1, 4, 2, 3
            )
            direction = maybe_encode_direction(
                self.posemb_version, position=position, direction=direction
            ).permute(0, 1, 4, 2, 3)
            position = encode_position(self.posemb_version, position=position).permute(
                0, 1, 4, 2, 3
            )
        x = torch.cat([channels, position, direction], dim=-3)
        *batch_shape, in_features, height, width = x.shape
        return (
            super()
            .forward(x.view(-1, in_features, height, width))
            .view(*batch_shape, -1, height, width)
        )

        :param channels: [batch_shape, inner_batch_shape, n_channels, height, width]
        :param position: [batch_shape, inner_batch_shape, 3, height, width]
        :param direction: [batch_shape, inner_batch_shape, 3, height, width]
        :return: [*batch_shape, out_features, height, width]
        
shap_e/models/nn/encoding.py