Project: s41467-024-47120-y | Function: forward_encoder (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** 

**KEY TERMS:** {}

**CONTEXT CODE:**         return kl_loss, recon_loss

    def encode(self, x):
        """   
        Do the VAE forward step to get the latent space of a tensor 

        :param x: list of tensors of longs, input sentence x
        :return: vector representing encoded latent space
        """
        if not isinstance(x,list):
            x = [x]
            
        # Encoder: x -> z, kl_loss
        z, kl_loss, mu = self.forward_encoder(x, return_mu=True)       

        return mu

    def decode(self, z, x=None):
        if x is not None:
            rl, y = self.forward_decoder(x,z,return_y=True)
            xr = y.argmax(2)
            # trim off eos
            #xr = [i[:(i == self.eos).nonzero()[0]] for i in xr]
            # trim off eos
            xrt = []
            for i in xr:
                try:
                    q = i[:(i==self.eos).nonzero()[0]]
                    xrt.append(q)
                except IndexError:
                    xrt.append(i)
            smiles = [self.tensor2string(i_x) for i_x in xrt]
        else:
            smiles = self.sample(z.shape[0], z=z, multinomial=True)
        return smiles

    def get_latent_space_from_smiles(self, smiles):
        """   
        Do the VAE forward step to get the latent space of a smiles 

        :param x: list of tensors of longs, input sentence x
        :return: vector representing encoded latent space
        """

        x = self.string2tensor(smiles)
        # Encoder: x -> z, kl_loss
        z, kl_loss = self.forward_encoder([x])       

        return z[0]

    def forward_encoder(self, x, return_mu=False):
        """Encoder step, emulating z ~ E(x) = q_E(z|x)

        :param x: list of tensors of longs, input sentence x
        :return: (n_batch, d_z) of floats, sample of latent vector z
        :return: float, kl term component of loss
        """
        # MASKED: VAE encoder forward pass that converts molecular SMILES to latent space encoding with KL divergence loss calculation
        raise NotImplementedError("This function is masked for testing purposes.")
	
    def forward_decoder(self, x, z, return_y=False):
        """Decoder step, emulating x ~ G(z)

        :param x: list of tensors of longs, input sentence x
        :param z: (n_batch, d_z) of floats, latent vector z
        :return: float, recon component of loss
        """
        # MASKED: VAE decoder forward pass that converts latent space back to molecular SMILES with reconstruction loss calculation
        raise NotImplementedError("This function is masked for testing purposes.")

    def sample_z_prior(self, n_batch):
        """Sampling z ~ p(z) = N(0, I)

        :param n_batch: number of batches
        :return: (n_batch, d_z) of floats, sample of latent z
        """

        return torch.randn(n_batch, self.q_mu.out_features,
                           device=self.x_emb.weight.device)

    def perturb_z(self, z, noise_norm, constant_norm=False):
        if noise_norm > 0.0:
            noise_vec = np.random.normal(0, 1, size=z.shape)
            noise_vec = noise_vec / np.linalg.norm(noise_vec)
            if constant_norm:
                return z + (noise_norm * noise_vec)
            else:
                noise_amp = np.random.uniform(
                    0, noise_norm, size=(z.shape[0], 1))
                return z + torch.tensor(noise_amp * noise_vec, dtype=z.dtype)
        else:
            return z

    def sample(self, n_batch, max_len=100, z=None, temp=1.0, multinomial=True):
        """Generating n_batch samples in eval mode (`z` could be
        not on same device)

        :param n_batch: number of sentences to generate
        :param max_len: max len of samples
        :param z: (n_batch, d_z) of floats, latent vector z or None
        :param temp: temperature of softmax
        :return: list of tensors of strings, samples sequence x
        """
        with torch.no_grad():
            if z is None:
                z = self.sample_z_prior(n_batch)
            z = z.to(self.device)
            z_0 = z.unsqueeze(1)

            # Initial values
            h = self.decoder_lat(z)
            h = h.unsqueeze(0).repeat(self.decoder_rnn.num_layers, 1, 1)
            w = torch.tensor(self.bos, device=self.device).repeat(n_batch)
            x = torch.tensor([self.pad], device=self.device).repeat(n_batch,
                                                                    max_len)
            x[:, 0] = self.bos
            end_pads = torch.tensor([max_len], device=self.device).repeat(
                n_batch)
            eos_mask = torch.zeros(n_batch, dtype=torch.bool,
                                   device=self.device)

            # Generating cycle
            for i in range(1, max_len):
                x_emb = self.x_emb(w).unsqueeze(1)
                x_input = torch.cat([x_emb, z_0], dim=-1)

                o, h = self.decoder_rnn(x_input, h)
                y = self.decoder_fc(o.squeeze(1))
                y = F.softmax(y / temp, dim=-1)
                if multinomial:
                    w = torch.multinomial(y, 1)[:, 0]
                else:
                    w = torch.argmax(y,1)


                x[~eos_mask, i] = w[~eos_mask]
                i_eos_mask = ~eos_mask & (w == self.eos)
                # new pytorch error with bool vs byte scaler 
                # try to convert it to byte tensor
                test_condition = torch.zeros((w.shape)).bool().to(self.device)
                test_condition = test_condition | (w==self.eos)
                #i_eos_mask = ~eos_mask & test_condition

                end_pads[i_eos_mask] = i + 1
                eos_mask = eos_mask | i_eos_mask

            # Converting `x` to list of tensors
            new_x = []
            for i in range(x.size(0)):
                new_x.append(x[i, :end_pads[i]])


Complete the function implementation replacing "raise NotImplementedError".