(x,
                                  n_token,
                                  d_embed,
                                  d_proj,
                                  cutoffs,
                                  initializer,
                                  proj_initializer,
                                  div_val=1,
                                  perms=None,
                                  proj_same_dim=True,
                                  scope='adaptive_embed')
return y, ret_params
perms: If None, first compute W = W1 x W2 (projection for each bin),
      and then compute X x W (embedding lookup). If not None,
      use bin-based embedding lookup with max_bin_size defined by
      the shape of perms.
perms: If None, first compute W = W1 x W2 (projection for each bin),
      and then compute X x W (embedding lookup). If not None,
      use bin-based embedding lookup with max_bin_size defined by
      the shape of perms.
perms : If None first compute W = W1 x W2 ( projection for each bin ) and then compute X x W ( embedding lookup ) . If not None use bin - based embedding lookup with max_bin_size defined by the shape of perms .
def mul_adaptive_embedding_lookup(x,
                                  n_token,
                                  d_embed,
                                  d_proj,
                                  cutoffs,
                                  initializer,
                                  proj_initializer,
                                  div_val=1,
                                  perms=None,
                                  proj_same_dim=True,
                                  scope='adaptive_embed'):
  """
  perms: If None, first compute W = W1 x W2 (projection for each bin),
      and then compute X x W (embedding lookup). If not None,
      use bin-based embedding lookup with max_bin_size defined by
      the shape of perms.
  """
  emb_scale = d_proj**0.5
  with tf.variable_scope(scope):
    if div_val == 1:
      lookup_table = tf.get_variable(
          'lookup_table', [n_token, d_embed], initializer=initializer)
      y = embedding_lookup(lookup_table, x)
      if d_proj != d_embed:
        proj_W = tf.get_variable(
            'proj_W', [d_embed, d_proj], initializer=proj_initializer)
        y = tf.einsum('ibe,ed->ibd', y, proj_W)
      else:
        proj_W = None
      ret_params = [lookup_table, proj_W]
    else:
      tables, projs = [], []
      cutoff_ends = [0] + cutoffs + [n_token]
      x_size = tf.shape(x)
      if perms is None:
        cat_lookup = []
      else:
        cat_lookup = tf.zeros([x_size[0], x_size[1], d_proj])
      for i in range(len(cutoff_ends) - 1):
        with tf.variable_scope('cutoff_{}'.format(i)):
          l_idx, r_idx = cutoff_ends[i], cutoff_ends[i + 1]
          cur_d_embed = d_embed // (div_val**i)
          lookup_table = tf.get_variable(
              'lookup_table', [r_idx - l_idx, cur_d_embed],
              initializer=initializer)
          if cur_d_embed == d_proj and not proj_same_dim:
            proj_W = None
          else:
            proj_W = tf.get_variable(
                'proj_W', [cur_d_embed, d_proj], initializer=proj_initializer)
          if perms is None:
            cat_lookup.append(tf.einsum('ie,ed->id', lookup_table, proj_W))
          else:
            # speed up the computation of the first bin
            # also save some meory
            if i == 0:
              cur_y = embedding_lookup(lookup_table, tf.minimum(x, r_idx - 1))
              if proj_W is not None:
                cur_y = tf.einsum('ibe,ed->ibd', cur_y, proj_W)
              cur_y *= perms[i][:, :, None]
              cat_lookup += cur_y
            else:
              cur_x = tf.einsum('ib,ibk->k', tf.to_float(x - l_idx), perms[i])
              cur_x = tf.to_int32(cur_x)
              cur_y = embedding_lookup(lookup_table, cur_x)
              if proj_W is not None:
                cur_y = tf.einsum('ke,ed->kd', cur_y, proj_W)
              cat_lookup += tf.einsum('kd,ibk->ibd', cur_y, perms[i])
          tables.append(lookup_table)
          projs.append(proj_W)
      if perms is None:
        cat_lookup = tf.concat(cat_lookup, 0)
        y = embedding_lookup(cat_lookup, x)
      else:
        y = cat_lookup
      ret_params = [tables, projs]

  y *= emb_scale
  return y, ret_params
def mul_adaptive_embedding_lookup ( x , n_token , d_embed , d_proj , cutoffs , initializer , proj_initializer , div_val = 1 , perms = None , proj_same_dim = True , scope = 'adaptive_embed' ) : emb_scale = d_proj ** 0.5 with tf . variable_scope ( scope ) : if div_val == 1 : lookup_table = tf . get_variable ( 'lookup_table' , [ n_token , d_embed ] , initializer = initializer ) y = embedding_lookup ( lookup_table , x ) if d_proj != d_embed : proj_W = tf . get_variable ( 'proj_W' , [ d_embed , d_proj ] , initializer = proj_initializer ) y = tf . einsum ( 'ibe,ed->ibd' , y , proj_W ) else : proj_W = None ret_params = [ lookup_table , proj_W ] else : tables , projs = [ ] , [ ] cutoff_ends = [ 0 ] + cutoffs + [ n_token ] x_size = tf . shape ( x ) if perms is None : cat_lookup = [ ] else : cat_lookup = tf . zeros ( [ x_size [ 0 ] , x_size [ 1 ] , d_proj ] ) for i in range ( len ( cutoff_ends ) - 1 ) : with tf . variable_scope ( 'cutoff_{}' . format ( i ) ) : l_idx , r_idx = cutoff_ends [ i ] , cutoff_ends [ i + 1 ] cur_d_embed = d_embed // ( div_val ** i ) lookup_table = tf . get_variable ( 'lookup_table' , [ r_idx - l_idx , cur_d_embed ] , initializer = initializer ) if cur_d_embed == d_proj and not proj_same_dim : proj_W = None else : proj_W = tf . get_variable ( 'proj_W' , [ cur_d_embed , d_proj ] , initializer = proj_initializer ) if perms is None : cat_lookup . append ( tf . einsum ( 'ie,ed->id' , lookup_table , proj_W ) ) else : # speed up the computation of the first bin # also save some meory if i == 0 : cur_y = embedding_lookup ( lookup_table , tf . minimum ( x , r_idx - 1 ) ) if proj_W is not None : cur_y = tf . einsum ( 'ibe,ed->ibd' , cur_y , proj_W ) cur_y *= perms [ i ] [ : , : , None ] cat_lookup += cur_y else : cur_x = tf . einsum ( 'ib,ibk->k' , tf . to_float ( x - l_idx ) , perms [ i ] ) cur_x = tf . to_int32 ( cur_x ) cur_y = embedding_lookup ( lookup_table , cur_x ) if proj_W is not None : cur_y = tf . einsum ( 'ke,ed->kd' , cur_y , proj_W ) cat_lookup += tf . einsum ( 'kd,ibk->ibd' , cur_y , perms [ i ] ) tables . append ( lookup_table ) projs . append ( proj_W ) if perms is None : cat_lookup = tf . concat ( cat_lookup , 0 ) y = embedding_lookup ( cat_lookup , x ) else : y = cat_lookup ret_params = [ tables , projs ] y *= emb_scale return y , ret_params