code: def top_k_hot(distances, k):
      """Returns the one-hot mask for k-nearest neighbours."""
      _, idx = jax.lax.top_k(-distances, k)
      tops = jax.nn.one_hot(idx, distances.shape[-1], dtype=jnp.int32)
      return tops.sum(axis=-2)
  
imports: from typing import Any
  import flax.linen
  from flax.linen.initializers import zeros
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.projects.performer import performer
label_desc: Use image node_id 2 for calling a global function with example usage: top_k_hot(distances, k) and returns: tops