class_docstr: ## Configurations
class_docstr_tok: ["#","#","Configurations"]
docstr: ### Sample $w$
  
  This samples $z$ randomly and get $w$ from the mapping network.
  
  We also apply style mixing sometimes where we generate two latent variables
  $z_1$ and $z_2$ and get corresponding $w_1$ and $w_2$.
  Then we randomly sample a cross-over point and apply $w_1$ to
  the generator blocks before the cross-over point and
  $w_2$ to the blocks after.
docstr_tok: ["#","#","#","Sample","$","w","$","This","samples","$","z","$","randomly","and","get","$","w","$","from","the","mapping","network",".","We","also","apply","style","mixing","sometimes","where","we","generate","two","latent","variables","$","z_1","$","and","$","z_2","$","and","get","corresponding","$","w_1","$","and","$","w_2","$",".","Then","we","randomly","sample","a","cross-over","point","and","apply","$","w_1","$","to","the","generator","blocks","before","the","cross-over","point","and","$","w_2","$","to","the","blocks","after","."]
code: def get_w(self, batch_size: int):
      """
      ### Sample $w$
  
      This samples $z$ randomly and get $w$ from the mapping network.
  
      We also apply style mixing sometimes where we generate two latent variables
      $z_1$ and $z_2$ and get corresponding $w_1$ and $w_2$.
      Then we randomly sample a cross-over point and apply $w_1$ to
      the generator blocks before the cross-over point and
      $w_2$ to the blocks after.
      """
  
      # Mix styles
      if torch.rand(()).item() < self.style_mixing_prob:
          # Random cross-over point
          cross_over_point = int(
              torch.rand(()).item() * self.n_gen_blocks
          )
          # Sample $z_1$ and $z_2$
          z2 = torch.randn(batch_size, self.d_latent).to(self.device)
          z1 = torch.randn(batch_size, self.d_latent).to(self.device)
          # Get $w_1$ and $w_2$
          w1 = self.mapping_network(z1)
          w2 = self.mapping_network(z2)
          # Expand $w_1$ and $w_2$ for the generator blocks and concatenate
          w1 = w1[None, :, :].expand(cross_over_point, -1, -1)
          w2 = w2[None, :, :].expand(
              self.n_gen_blocks - cross_over_point, -1, -1
          )
          return torch.cat((w1, w2), dim=0)
      # Without mixing
      else:
          # Sample $z$ and $z$
          z = torch.randn(batch_size, self.d_latent).to(self.device)
          # Get $w$ and $w$
          w = self.mapping_network(z)
          # Expand $w$ for the generator blocks
          return w[None, :, :].expand(self.n_gen_blocks, -1, -1)
  
imports: import math
  from pathlib import Path
  from typing import Iterator, Tuple
  import torch
  import torch.utils.data
  import torchvision
  from PIL import Image
  from labml import tracker, lab, monit, experiment
  from labml.configs import BaseConfigs
  from labml_helpers.device import DeviceConfigs
  from labml_helpers.train_valid import ModeState, hook_model_outputs
  from labml_nn.gan.stylegan import Discriminator, Generator, MappingNetwork, GradientPenalty, PathLengthPenalty
  from labml_nn.gan.wasserstein import DiscriminatorLoss, GeneratorLoss
  from labml_nn.utils import cycle_dataloader
label_desc: Use image node_id 2 for calling the Configs obj's underlying member method code with example usage: obj.get_w(batch_size) and returns: torch, w