def translate()

in latent_operators.py [0:0]

• 8 lines of code
• 4 McCabe index (conditional complexity)

    def translate(self, z, shift):
        """Translate latent

        Args:
            z (1-dim tensor): latent vector
            shift (int): amount by which to shift.
                shift of 0 corresponds to the identity.
        """
        # reshape into 2D tensor
        z_2d = z.reshape(self.n_rotations + 1, -1)
        translation_matrix = self.translation_matrices[shift]
        # move to cpu if tensor is cpu. Used for eval
        if not z_2d.is_cuda:
            translation_matrix = translation_matrix.cpu()
        # translation
        z_2d_shifted = translation_matrix.matmul(z_2d)
        # reshape back
        z_shifted = z_2d_shifted.reshape(z.shape)
        return z_shifted