def unpatchify(self, patchified_pixel_values):
        """
        Args:
            patchified_pixel_values (`tf.Tensor` of shape `(batch_size, num_patches, patch_size**2 * num_channels)`:
                Patchified pixel values.

        Returns:
            `tf.Tensor` of shape `(batch_size, height, width, num_channels)`:
                Pixel values.
        """
        patch_size, num_channels = self.config.patch_size, self.config.num_channels
        num_patches_one_direction = int(shape_list(patchified_pixel_values)[1] ** 0.5)
        # sanity check
        tf.debugging.assert_equal(
            num_patches_one_direction * num_patches_one_direction,
            shape_list(patchified_pixel_values)[1],
            message="Make sure that the number of patches can be squared",
        )

        # unpatchify
        batch_size = shape_list(patchified_pixel_values)[0]
        patchified_pixel_values = tf.reshape(
            patchified_pixel_values,
            (batch_size, num_patches_one_direction, num_patches_one_direction, patch_size, patch_size, num_channels),
        )
        patchified_pixel_values = tf.einsum("nhwpqc->nhpwqc", patchified_pixel_values)
        pixel_values = tf.reshape(
            patchified_pixel_values,
            (batch_size, num_patches_one_direction * patch_size, num_patches_one_direction * patch_size, num_channels),
        )
        return pixel_values

        Args:
            patchified_pixel_values (`tf.Tensor` of shape `(batch_size, num_patches, patch_size**2 * num_channels)`:
                Patchified pixel values.

        Returns:
            `tf.Tensor` of shape `(batch_size, height, width, num_channels)`:
                Pixel values.
        
dpkd/transformers/src/transformers/models/vit_mae/modeling_tf_vit_mae.py