original_docstring:  Build a map from TF variables to PyTorch modules. 
docstring: Build a map from TF variables to PyTorch modules.
docstring_tokens: Build a map from TF variables to PyTorch modules .
code: def build_tf_to_pytorch_map(model, config):
      """ Build a map from TF variables to PyTorch modules. """
      tf_to_pt_map = {}
  
      # Embeddings and GenZ
      tf_to_pt_map.update({'linear/w/ema_0.9999': model.embeddings.weight,
                           'Generator/GenZ/G_linear/b/ema_0.9999': model.generator.gen_z.bias,
                           'Generator/GenZ/G_linear/w/ema_0.9999': model.generator.gen_z.weight_orig,
                           'Generator/GenZ/G_linear/u0': model.generator.gen_z.weight_u})
  
      # GBlock blocks
      model_layer_idx = 0
      for i, (up, in_channels, out_channels) in enumerate(config.layers):
          if i == config.attention_layer_position:
              model_layer_idx += 1
          layer_str = "Generator/GBlock_%d/" % i if i > 0 else "Generator/GBlock/"
          layer_pnt = model.generator.layers[model_layer_idx]
          for i in range(4):  #  Batchnorms
              batch_str = layer_str + ("BatchNorm_%d/" % i if i > 0 else "BatchNorm/")
              batch_pnt = getattr(layer_pnt, 'bn_%d' % i)
              for name in ('offset', 'scale'):
                  sub_module_str = batch_str + name + "/"
                  sub_module_pnt = getattr(batch_pnt, name)
                  tf_to_pt_map.update({sub_module_str + "w/ema_0.9999": sub_module_pnt.weight_orig,
                                       sub_module_str + "u0": sub_module_pnt.weight_u})
          for i in range(4):  # Convolutions
              conv_str = layer_str + "conv%d/" % i
              conv_pnt = getattr(layer_pnt, 'conv_%d' % i)
              tf_to_pt_map.update({conv_str + "b/ema_0.9999": conv_pnt.bias,
                                   conv_str + "w/ema_0.9999": conv_pnt.weight_orig,
                                   conv_str + "u0": conv_pnt.weight_u})
          model_layer_idx += 1
  
      # Attention block
      layer_str = "Generator/attention/"
      layer_pnt = model.generator.layers[config.attention_layer_position]
      tf_to_pt_map.update({layer_str + "gamma/ema_0.9999": layer_pnt.gamma})
      for pt_name, tf_name in zip(['snconv1x1_g', 'snconv1x1_o_conv', 'snconv1x1_phi', 'snconv1x1_theta'],
                                  ['g/', 'o_conv/', 'phi/', 'theta/']):
          sub_module_str = layer_str + tf_name
          sub_module_pnt = getattr(layer_pnt, pt_name)
          tf_to_pt_map.update({sub_module_str + "w/ema_0.9999": sub_module_pnt.weight_orig,
                               sub_module_str + "u0": sub_module_pnt.weight_u})
  
      # final batch norm and conv to rgb
      layer_str = "Generator/BatchNorm/"
      layer_pnt = model.generator.bn
      tf_to_pt_map.update({layer_str + "offset/ema_0.9999": layer_pnt.bias,
                           layer_str + "scale/ema_0.9999": layer_pnt.weight})
      layer_str = "Generator/conv_to_rgb/"
      layer_pnt = model.generator.conv_to_rgb
      tf_to_pt_map.update({layer_str + "b/ema_0.9999": layer_pnt.bias,
                           layer_str + "w/ema_0.9999": layer_pnt.weight_orig,
                           layer_str + "u0": layer_pnt.weight_u})
      return tf_to_pt_map
code_tokens: def build_tf_to_pytorch_map ( model , config ) : tf_to_pt_map = { } tf_to_pt_map . update ( { 'linear/w/ema_0.9999' : model . embeddings . weight , 'Generator/GenZ/G_linear/b/ema_0.9999' : model . generator . gen_z . bias , 'Generator/GenZ/G_linear/w/ema_0.9999' : model . generator . gen_z . weight_orig , 'Generator/GenZ/G_linear/u0' : model . generator . gen_z . weight_u } ) model_layer_idx = 0 for i , ( up , in_channels , out_channels ) in enumerate ( config . layers ) : if i == config . attention_layer_position : model_layer_idx += 1 layer_str = "Generator/GBlock_%d/" % i if i > 0 else "Generator/GBlock/" layer_pnt = model . generator . layers [ model_layer_idx ] for i in range ( 4 ) : batch_str = layer_str + ( "BatchNorm_%d/" % i if i > 0 else "BatchNorm/" ) batch_pnt = getattr ( layer_pnt , 'bn_%d' % i ) for name in ( 'offset' , 'scale' ) : sub_module_str = batch_str + name + "/" sub_module_pnt = getattr ( batch_pnt , name ) tf_to_pt_map . update ( { sub_module_str + "w/ema_0.9999" : sub_module_pnt . weight_orig , sub_module_str + "u0" : sub_module_pnt . weight_u } ) for i in range ( 4 ) : conv_str = layer_str + "conv%d/" % i conv_pnt = getattr ( layer_pnt , 'conv_%d' % i ) tf_to_pt_map . update ( { conv_str + "b/ema_0.9999" : conv_pnt . bias , conv_str + "w/ema_0.9999" : conv_pnt . weight_orig , conv_str + "u0" : conv_pnt . weight_u } ) model_layer_idx += 1 layer_str = "Generator/attention/" layer_pnt = model . generator . layers [ config . attention_layer_position ] tf_to_pt_map . update ( { layer_str + "gamma/ema_0.9999" : layer_pnt . gamma } ) for pt_name , tf_name in zip ( [ 'snconv1x1_g' , 'snconv1x1_o_conv' , 'snconv1x1_phi' , 'snconv1x1_theta' ] , [ 'g/' , 'o_conv/' , 'phi/' , 'theta/' ] ) : sub_module_str = layer_str + tf_name sub_module_pnt = getattr ( layer_pnt , pt_name ) tf_to_pt_map . update ( { sub_module_str + "w/ema_0.9999" : sub_module_pnt . weight_orig , sub_module_str + "u0" : sub_module_pnt . weight_u } ) layer_str = "Generator/BatchNorm/" layer_pnt = model . generator . bn tf_to_pt_map . update ( { layer_str + "offset/ema_0.9999" : layer_pnt . bias , layer_str + "scale/ema_0.9999" : layer_pnt . weight } ) layer_str = "Generator/conv_to_rgb/" layer_pnt = model . generator . conv_to_rgb tf_to_pt_map . update ( { layer_str + "b/ema_0.9999" : layer_pnt . bias , layer_str + "w/ema_0.9999" : layer_pnt . weight_orig , layer_str + "u0" : layer_pnt . weight_u } ) return tf_to_pt_map
short_docstring: Build a map from TF variables to PyTorch modules.
short_docstring_tokens: Build a map from TF variables to PyTorch modules .
comment: """ Build a map from TF variables to PyTorch modules. """$SEP$# Embeddings and GenZ$SEP$# GBlock blocks$SEP$#  Batchnorms$SEP$# Convolutions$SEP$# Attention block$SEP$# final batch norm and conv to rgb