text: import torch
  import torch.nn as nn
  from torch.nn import init
  import torch.nn.functional as F
  import math
  
  import numpy as np
  
  torch.set_printoptions(precision=2,threshold=float('inf'))
  
  class AGCNBlock(nn.Module):
      def __init__(self,input_dim,hidden_dim,gcn_layer=2,dropout=0.0,relu=0):
          super(AGCNBlock,self).__init__()
          if dropout > 0.001:
              self.dropout_layer = nn.Dropout(p=dropout)
          self.sort = 'sort'
          self.model='agcn'
          self.gcns=nn.ModuleList()
          self.bn = 0
          self.add_self = 1
          self.normalize_embedding = 1
          self.gcns.append(GCNBlock(input_dim,hidden_dim,self.bn,self.add_self,self.normalize_embedding,dropout,relu))
          self.pool = 'mean'
          self.tau = 1.
          self.lamda = 1.
  
          for i in range(gcn_layer-1):
              if i==gcn_layer-2 and (not 1):
                  self.gcns.append(GCNBlock(hidden_dim,hidden_dim,self.bn,self.add_self,self.normalize_embedding,dropout,0))
              else:
                  self.gcns.append(GCNBlock(hidden_dim,hidden_dim,self.bn,self.add_self,self.normalize_embedding,dropout,relu))
              
          if self.model=='diffpool':
              self.pool_gcns=nn.ModuleList()
              tmp=input_dim
              self.diffpool_k=200
              for i in range(3):
                  self.pool_gcns.append(GCNBlock(tmp,200,0,0,0,dropout,relu))
                  tmp=200
  
          self.w_a=nn.Parameter(torch.zeros(1,hidden_dim,1))
          self.w_b=nn.Parameter(torch.zeros(1,hidden_dim,1))
          torch.nn.init.normal_(self.w_a)
          torch.nn.init.uniform_(self.w_b,-1,1)
  
          self.pass_dim=hidden_dim
  
          if self.pool=='mean':
              self.pool=self.mean_pool
          elif self.pool=='max':
              self.pool=self.max_pool
          elif self.pool=='sum':
              self.pool=self.sum_pool
  
          self.softmax='global'
          if self.softmax=='gcn':
              self.att_gcn=GCNBlock(2,1,0,0,dropout,relu)
          self.khop=1
          self.adj_norm='none'
  
          self.filt_percent=0.25       #default 0.5
          self.eps=1e-10
  
          self.tau_config=1
          if 1==-1.:
              self.tau=nn.Parameter(torch.tensor(1),requires_grad=False)
          elif 1==-2.:
              self.tau_fc=nn.Linear(hidden_dim,1)
              torch.nn.init.constant_(self.tau_fc.bias,1)
              torch.nn.init.xavier_normal_(self.tau_fc.weight.t())
          else:
              self.tau=nn.Parameter(torch.tensor(self.tau))
          self.lamda1=nn.Parameter(torch.tensor(self.lamda))
          self.lamda2=nn.Parameter(torch.tensor(self.lamda))
  
          self.att_norm=0
          
          self.dnorm=0
          self.dnorm_coe=1
  
          self.att_out=0
          self.single_att=0
  
  
      def forward(self,X,adj,mask,is_print=False):
          '''
      input:
          X:  node input features , [batch,node_num,input_dim],dtype=float
          adj: adj matrix, [batch,node_num,node_num], dtype=float
          mask: mask for nodes, [batch,node_num]
      outputs:
          out:unormalized classification prob, [batch,hidden_dim]
          H: batch of node hidden features, [batch,node_num,pass_dim]
          new_adj: pooled new adj matrix, [batch, k_max, k_max]
          new_mask: [batch, k_max]
          '''
          hidden=X
          #adj = adj.float()
          # print('input size:')
          # print(hidden.shape)
          
          is_print1=is_print2=is_print
          if adj.shape[-1]>100:
              is_print1=False
  
          for gcn in self.gcns:
              hidden=gcn(hidden,adj,mask)
          #     print('gcn:')
          #     print(hidden.shape)
          # print('mask:')
          # print(mask.unsqueeze(2).shape)
          # print(mask.sum(dim=1))
  
          hidden=mask.unsqueeze(2)*hidden
          # print(hidden[0][0])
          # print(hidden[0][-1])
  
          if self.model=='unet':
              att=torch.matmul(hidden,self.w_a).squeeze()
              att=att/torch.sqrt((self.w_a.squeeze(2)**2).sum(dim=1,keepdim=True))
          elif self.model=='agcn':
              if self.softmax=='global' or self.softmax=='mix':
                  if False:
                      dgree_w = torch.sum(adj, dim=2) / torch.sum(adj, dim=2).max(1, keepdim=True)[0]
                      att_a=torch.matmul(hidden,self.w_a).squeeze()*dgree_w+(mask-1)*1e10
                  else:
                      att_a=torch.matmul(hidden,self.w_a).squeeze()+(mask-1)*1e10
                      # print(att_a[0][:10])
                      # print(att_a[0][-10:-1])
                  att_a_1=att_a=torch.nn.functional.softmax(att_a,dim=1)
                  # print(att_a[0][:10])
                  # print(att_a[0][-10:-1])
  
                  if self.dnorm:
                      scale=mask.sum(dim=1,keepdim=True)/self.dnorm_coe
                      att_a=scale*att_a
              if self.softmax=='neibor' or self.softmax=='mix':
                  att_b=torch.matmul(hidden,self.w_b).squeeze()+(mask-1)*1e10
                  att_b_max,_=att_b.max(dim=1,keepdim=True)
                  if self.tau_config!=-2:
                      att_b=torch.exp((att_b-att_b_max)*torch.abs(self.tau))
                  else:
                      att_b=torch.exp((att_b-att_b_max)*torch.abs(self.tau_fc(self.pool(hidden,mask))))
                  denom=att_b.unsqueeze(2)
                  for _ in range(self.khop):
                      denom=torch.matmul(adj,denom)
                  denom=denom.squeeze()+self.eps
                  att_b=(att_b*torch.diagonal(adj,0,1,2))/denom
                  if self.dnorm:
                      if self.adj_norm=='diag':
                          diag_scale=mask/(torch.diagonal(adj,0,1,2)+self.eps)
                      elif self.adj_norm=='none':
                          diag_scale=adj.sum(dim=1)
                      att_b=att_b*diag_scale
                  att_b=att_b*mask
                          
              if self.softmax=='global':
                  att=att_a
              elif self.softmax=='neibor' or self.softmax=='hardnei':
                  att=att_b
              elif self.softmax=='mix':
                  att=att_a*torch.abs(self.lamda1)+att_b*torch.abs(self.lamda2)
          # print('att:')
          # print(att.shape)
          Z=hidden
  
          if self.model=='unet':
              Z=torch.tanh(att.unsqueeze(2))*Z
          elif self.model=='agcn':
              if self.single_att:
                  Z=Z
              else:
                  Z=att.unsqueeze(2)*Z
          # print('Z shape')
          # print(Z.shape)
          k_max=int(math.ceil(self.filt_percent*adj.shape[-1]))
          # print('k_max')
          # print(k_max)
          if self.model=='diffpool':
              k_max=min(k_max,self.diffpool_k)
  
          k_list=[int(math.ceil(self.filt_percent*x)) for x in mask.sum(dim=1).tolist()]
          # print('k_list')
          # print(k_list)
          if self.model!='diffpool': 
              if self.sort=='sample':
                  att_samp = att * mask
                  att_samp = (att_samp/att_samp.sum(1)).detach().cpu().numpy()
                  top_index = ()
                  for i in range(att.size(0)):
                      top_index = (torch.LongTensor(np.random.choice(att_samp.size(1), k_max, att_samp[i])) ,)
                  top_index = torch.stack(top_index,1)
              elif self.sort=='random_sample':
                  top_index = torch.LongTensor(att.size(0), k_max)*0
                  for i in range(att.size(0)):
                      top_index[i,0:k_list[i]] = torch.randperm(int(mask[i].sum().item()))[0:k_list[i]] 
              else: #sort
                  _,top_index=torch.topk(att,k_max,dim=1)
          # print('top_index')
          # print(top_index)
          # print(len(top_index[0]))
          new_mask=X.new_zeros(X.shape[0],k_max)
          # print('new_mask')
          # print(new_mask.shape)
          visualize_tools=None 
          if self.model=='unet':
              for i,k in enumerate(k_list):
                  for j in range(int(k),k_max):
                      top_index[i][j]=adj.shape[-1]-1
                      new_mask[i][j]=-1.
              new_mask=new_mask+1
              top_index,_=torch.sort(top_index,dim=1)
              assign_m=X.new_zeros(X.shape[0],k_max,adj.shape[-1])
              for i,x in enumerate(top_index):
                  assign_m[i]=torch.index_select(adj[i],0,x)
              new_adj=X.new_zeros(X.shape[0],k_max,k_max)
              H=Z.new_zeros(Z.shape[0],k_max,Z.shape[-1])
              for i,x in enumerate(top_index):
                  new_adj[i]=torch.index_select(assign_m[i],1,x)
                  H[i]=torch.index_select(Z[i],0,x)
  
          elif self.model=='agcn':
              assign_m=X.new_zeros(X.shape[0],k_max,adj.shape[-1])
              # print('assign_m.shape')
              # print(assign_m.shape)
              for i,k in enumerate(k_list):
                  #print('top_index[i][j]')
                  for j in range(int(k)):  
                      #print(str(top_index[i][j].item())+' ', end='')
                      assign_m[i][j]=adj[i][top_index[i][j]]
                      #print(assign_m[i][j])
                      new_mask[i][j]=1.
  
              assign_m=assign_m/(assign_m.sum(dim=1,keepdim=True)+self.eps)
              H=torch.matmul(assign_m,Z)
              # print('H')
              # print(H.shape)
              new_adj=torch.matmul(torch.matmul(assign_m,adj),torch.transpose(assign_m,1,2))
              # print(torch.matmul(assign_m,adj).shape)
              # print('new_adj:')
              # print(new_adj.shape)
              
          elif self.model=='diffpool':
              hidden1=X
              for gcn in self.pool_gcns:
                  hidden1=gcn(hidden1,adj,mask)
              assign_m=X.new_ones(X.shape[0],X.shape[1],k_max)*(-100000000.)
              for i,x in enumerate(hidden1):
                  k=min(k_list[i],k_max)
                  assign_m[i,:,0:k]=hidden1[i,:,0:k]
                  for j in range(int(k)):
                      new_mask[i][j]=1.
  
              assign_m=torch.nn.functional.softmax(assign_m,dim=2)*mask.unsqueeze(2)
              assign_m_t=torch.transpose(assign_m,1,2)
              new_adj=torch.matmul(torch.matmul(assign_m_t,adj),assign_m)
              H=torch.matmul(assign_m_t,Z)
          # print('pool')    
          if self.att_out and self.model=='agcn':
              if self.softmax=='global':
                  out=self.pool(att_a_1.unsqueeze(2)*hidden,mask)
              elif self.softmax=='neibor':
                  att_b_sum=att_b.sum(dim=1,keepdim=True)
                  out=self.pool((att_b/(att_b_sum+self.eps)).unsqueeze(2)*hidden,mask)
          else:
              # print('hidden.shape')
              # print(hidden.shape)
              out=self.pool(hidden,mask)
              # print('out shape')
              # print(out.shape)
             
          if self.adj_norm=='tanh' or self.adj_norm=='mix':
              new_adj=torch.tanh(new_adj)
          elif self.adj_norm=='diag' or self.adj_norm=='mix':
              diag_elem=torch.pow(new_adj.sum(dim=2)+self.eps,-0.5)
              diag=new_adj.new_zeros(new_adj.shape)
              for i,x in enumerate(diag_elem):
                  diag[i]=torch.diagflat(x)
              new_adj=torch.matmul(torch.matmul(diag,new_adj),diag)
  
          visualize_tools=[]
          '''
          if (not self.training) and is_print1:
              print('**********************************')
              print('node_feat:',X.type(),X.shape)
              print(X)
              if self.model!='diffpool':
                  print('**********************************')
                  print('att:',att.type(),att.shape)
                  print(att)
                  print('**********************************')
                  print('top_index:',top_index.type(),top_index.shape)
                  print(top_index)
              print('**********************************')
              print('adj:',adj.type(),adj.shape)
              print(adj)
              print('**********************************')
              print('assign_m:',assign_m.type(),assign_m.shape)
              print(assign_m)
              print('**********************************')
              print('new_adj:',new_adj.type(),new_adj.shape)
              print(new_adj)
              print('**********************************')
              print('new_mask:',new_mask.type(),new_mask.shape)
              print(new_mask)
          '''
          #visualization
          from os import path
          if not path.exists('att_1.pt'):
              torch.save(att[0], 'att_1.pt')
              torch.save(top_index[0], 'att_ind1.pt')
          elif not path.exists('att_2.pt'):
              torch.save(att[0], 'att_2.pt')
              torch.save(top_index[0], 'att_ind2.pt')
          else:
              torch.save(att[0], 'att_3.pt')
              torch.save(top_index[0], 'att_ind3.pt')
  
          if (not self.training) and is_print2:
              if self.model!='diffpool':
                  visualize_tools.append(att[0])
                  visualize_tools.append(top_index[0])
              visualize_tools.append(new_adj[0])
              visualize_tools.append(new_mask.sum())
          # print('**********************************')
          return out,H,new_adj,new_mask,visualize_tools
      
      def mean_pool(self,x,mask):
          return x.sum(dim=1)/(self.eps+mask.sum(dim=1,keepdim=True))
      
      def sum_pool(self,x,mask):
          return x.sum(dim=1)
  
      @staticmethod
      def max_pool(x,mask):
          #output: [batch,x.shape[2]]
          m=(mask-1)*1e10
          r,_=(x+m.unsqueeze(2)).max(dim=1)
          return r
  # GCN basic operation
  class GCNBlock(nn.Module):
      def __init__(self, input_dim, output_dim, bn=0,add_self=0, normalize_embedding=0,
              dropout=0.0,relu=0, bias=True):
          super(GCNBlock,self).__init__()
          self.add_self = add_self
          self.dropout = dropout
          self.relu=relu
          self.bn=bn
          if dropout > 0.001:
              self.dropout_layer = nn.Dropout(p=dropout)
          if self.bn:
              self.bn_layer = torch.nn.BatchNorm1d(output_dim)
  
          self.normalize_embedding = normalize_embedding
          self.input_dim = input_dim
          self.output_dim = output_dim
          self.weight = nn.Parameter(torch.FloatTensor(input_dim, output_dim).cuda())
          torch.nn.init.xavier_normal_(self.weight)
          if bias:
              self.bias = nn.Parameter(torch.zeros(output_dim).cuda())
          else:
              self.bias = None
  
      def forward(self, x, adj, mask):
          y = torch.matmul(adj, x)
          if self.add_self:
              y += x
          y = torch.matmul(y,self.weight)
          if self.bias is not None:
              y = y + self.bias
          if self.normalize_embedding:
              y = F.normalize(y, p=2, dim=2)
          if self.bn:
              index=mask.sum(dim=1).long().tolist()
              bn_tensor_bf=mask.new_zeros((sum(index),y.shape[2]))
              bn_tensor_af=mask.new_zeros(*y.shape)
              start_index=[]
              ssum=0
              for i in range(x.shape[0]):
                  start_index.append(ssum)
                  ssum+=index[i]
              start_index.append(ssum)
              for i in range(x.shape[0]):
                  bn_tensor_bf[start_index[i]:start_index[i+1]]=y[i,0:index[i]]
              bn_tensor_bf=self.bn_layer(bn_tensor_bf)
              for i in range(x.shape[0]):
                  bn_tensor_af[i,0:index[i]]=bn_tensor_bf[start_index[i]:start_index[i+1]]
              y=bn_tensor_af
          if self.dropout > 0.001:
              y = self.dropout_layer(y)
          if self.relu=='relu':
              y=torch.nn.functional.relu(y)
              print('hahah')
          elif self.relu=='lrelu':
              y=torch.nn.functional.leaky_relu(y,0.1)
          return y
  
  #experimental function, untested
  class masked_batchnorm(nn.Module):
      def __init__(self,feat_dim,epsilon=1e-10):
          super().__init__()
          self.alpha=nn.Parameter(torch.ones(feat_dim))
          self.beta=nn.Parameter(torch.zeros(feat_dim))
          self.eps=epsilon
  
      def forward(self,x,mask):
          '''
          x: node feat, [batch,node_num,feat_dim]
          mask: [batch,node_num]
          '''
          mask1 = mask.unsqueeze(2)
          mask_sum = mask.sum()
          mean = x.sum(dim=(0,1),keepdim=True)/(self.eps+mask_sum)
          temp = (x - mean)**2
          temp = temp*mask1
          var = temp.sum(dim=(0,1),keepdim=True)/(self.eps+mask_sum)
          rstd = torch.rsqrt(var+self.eps)
          x=(x-mean)*rstd
          return ((x*self.alpha) + self.beta)*mask1