code: def forward(self, dist, angle, idx_kj):
      dist = dist / self.cutoff
      rbf = torch.stack([f(dist) for f in self.bessel_funcs], dim=1)
      rbf = self.envelope(dist).unsqueeze(-1) * rbf
  
      cbf = torch.stack([f(angle) for f in self.sph_funcs], dim=1)
  
      n, k = self.num_spherical, self.num_radial
      out = (rbf[idx_kj].view(-1, n, k) * cbf.view(-1, n, 1)).view(
          -1, n * k
      )
      return out
  
imports: import numpy
  import torch
  from scipy.optimize import brentq
  from scipy import special
  from math import pi
label_desc: Use image node_id 2 for calling the angle_emb obj's underlying member method code with example usage: obj.forward(dist, angle, idx_kj) and returns: out