code: def _variance_estimator_dense(r, pos, f):
      nei_idx = (
          tp.ball_query(r, _MAX_NEIGHBOURS, pos, pos, sort=True)[0]
          .reshape(pos.shape[0], -1)
          .long()
      )  # [B,N * nei]
      f_neighboors = f.gather(1, nei_idx).reshape(
          f.shape[0], f.shape[1], -1
      )  # [B,N , nei]
      gradient = (
          f.unsqueeze(-1).repeat(1, 1, f_neighboors.shape[-1])
          - f_neighboors
      ) ** 2  # [B,N,nei]
      return gradient.sum(-1)
  
file_path: torch-points3d/torch_points3d/core/losses/dirichlet_loss.py
imports: import torch
  import torch_points_kernels
  from torch_cluster import radius
  from torch_scatter import scatter_add
label_desc: Use image node_id 3 for calling a global function with example usage: _variance_estimator_dense(r, pos, f) and returns: gradient