code: def one_hot(
      x, n_categories, dtype=jnp.float32
  ):  # pylint: disable=invalid-name
      """Makes a one-hot array (n+1 dims) from an int-categorical array (n dims)."""
      indices_less_than_n = jnp.arange(n_categories)
      return jnp.array(
          x[..., jnp.newaxis] == indices_less_than_n, dtype
      )
  
imports: from absl import logging
  import numpy
  import tensorflow
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
label_desc: Use image node_id 15 for calling a global function with example usage: one_hot(x, n_categories, dtype) and returns: jnp