text: # -*- coding: utf-8 -*-
  
  """Source code for simple multi-perceptron layer class.
  
  Author: Yoshinari Motokawa <yoshinari.moto@fuji.waseda.jp>
  """
  
  import torch.nn as nn
  from .vit import Block, Attention
  
  
  class HardShrinkBlock(Block):
      def __init__(
          self,
          dim,
          num_heads,
          mlp_ratio,
          qkv_bias=False,
          qk_scale=None,
          drop=0.0,
          attn_drop=0.0,
          drop_path=0.0,
          act_layer=nn.GELU,
          norm_layer=nn.LayerNorm,
          **kwargs
      ):
          super().__init__(
              dim,
              num_heads,
              mlp_ratio,
              qkv_bias,
              qk_scale,
              drop,
              attn_drop,
              drop_path,
              act_layer,
              norm_layer,
              **kwargs
          )
          self.attn = HardShrinkAttention(
              dim,
              num_heads=num_heads,
              qkv_bias=qkv_bias,
              qk_scale=qk_scale,
              attn_drop=attn_drop,
              proj_drop=drop,
              lambd=kwargs["af_lambd"],
          )
  
  
  class HardShrinkAttention(Attention):
      def __init__(
          self,
          dim,
          num_heads=8,
          qkv_bias=False,
          qk_scale=None,
          attn_drop=0.0,
          proj_drop=0.0,
          lambd=0.0,
      ):
          super().__init__(dim, num_heads, qkv_bias, qk_scale, attn_drop, proj_drop)
          self.hard_shrink = nn.Hardshrink(lambd=lambd)
  
      def forward(self, x):
          B, N, C = x.shape
          qkv = (
              self.qkv(x)
              .reshape(B, N, 3, self.num_heads, C // self.num_heads)
              .permute(2, 0, 3, 1, 4)
          )
          q, k, v = qkv[0], qkv[1], qkv[2]
  
          attn = (q @ k.transpose(-2, -1)) * self.scale
          attn = attn.softmax(dim=-1)
          attn = self.hard_shrink(attn)
          attn = self.attn_drop(attn)
  
          x = (attn @ v).transpose(1, 2).reshape(B, N, C)
          x = self.proj(x)
          x = self.proj_drop(x)
          return x
  
      def forward_attn(self, x):
          B, N, C = x.shape
          qkv = (
              self.qkv(x)
              .reshape(B, N, 3, self.num_heads, C // self.num_heads)
              .permute(2, 0, 3, 1, 4)
          )
          q, k, v = qkv[0], qkv[1], qkv[2]
  
          attn = (q @ k.transpose(-2, -1)) * self.scale
          attn = attn.softmax(dim=-1)
          attn = self.hard_shrink(attn)
          attn = self.attn_drop(attn)
  
          x = (attn @ v).transpose(1, 2).reshape(B, N, C)
          x = self.proj(x)
          x = self.proj_drop(x)
          return x, attn
  