code: def forward(self, x):
      mask = self.smax(self.weight)
      masked_x = torch.einsum("nd,bd->bnd", mask, x)  # [B, k, D]
      return masked_x
  
file_path: pytorch_tabular/src/pytorch_tabular/models/danet/arch_blocks.py
imports: from functools import partial
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from pytorch_tabular.models.common.layers.activations import entmax15
  from pytorch_tabular.models.common.layers.batch_norm import GBN
label_desc: Use image node_id 2 for calling the LearnableLocality obj's underlying member method code with example usage: obj.forward(x) and returns: masked_x