code: def _one_hot(x, k, dtype=np.float32):
      """Create a one-hot encoding of x of size k."""
      return np.array(x[:, None] == np.arange(k), dtype)
  
file_path: neural-tangents/examples/datasets.py
imports: import gzip
  import os
  import shutil
  import urllib.request
  from jax import random
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: _one_hot(x, k, dtype) and returns: np