def __init__(self, nx, n_ctx, config, scale=False, **kwargs):
    super(TFAttention, self).__init__(**kwargs)
    self.output_attentions = config.output_attentions
    n_state = nx
    assert ((n_state % config.n_head) == 0)
    self.n_ctx = n_ctx
    self.n_head = config.n_head
    self.split_size = n_state
    self.scale = scale
    self.c_attn = TFConv1D((n_state * 3), nx, initializer_range=config.initializer_range, name='c_attn')
    self.c_proj = TFConv1D(n_state, nx, initializer_range=config.initializer_range, name='c_proj')
    self.attn_dropout = tf.keras.layers.Dropout(config.attn_pdrop)
    self.resid_dropout = tf.keras.layers.Dropout(config.resid_pdrop)
    self.pruned_heads = set()
def __init__(self, nx, n_ctx, config, scale=False, **kwargs):
    super(TFAttention, self).__init__(**kwargs)
    self.output_attentions = config.output_attentions
    n_state = nx
    assert ((n_state % config.n_head) == 0)
    self.n_ctx = n_ctx
    self.n_head = config.n_head
    self.split_size = n_state
    self.scale = scale
    self.c_attn = TFConv1D((n_state * 3), nx, initializer_range=config.initializer_range, name='c_attn')
    self.c_proj = TFConv1D(n_state, nx, initializer_range=config.initializer_range, name='c_proj')
    self.attn_dropout = tf.keras.layers.Dropout(config.attn_pdrop)
    self.resid_dropout = tf.keras.layers.Dropout(config.resid_pdrop)
    self.pruned_heads = set()