code: def transform_to_2hot_nonlinear(scalar: Array, bins: Array) -> Array:
      """Transforms a scalar tensor to a 2 hot representation defined using bins."""
      min_value, max_value = bins[0], bins[-1]
      num_bins = len(bins)
      scalar = jnp.clip(scalar, min_value, max_value)
      upper_index = jnp.argmax(
          scalar[..., None] <= bins.reshape(1, -1), axis=-1
      )
      upper_value = bins[upper_index]
      lower_value = bins[upper_index - 1]
      p_lower = (upper_value - scalar) / (upper_value - lower_value)
      p_upper = 1 - p_lower
      lower_one_hot = (
          jax.nn.one_hot(upper_index - 1, num_bins) * p_lower[..., None]
      )
      upper_one_hot = (
          jax.nn.one_hot(upper_index, num_bins) * p_upper[..., None]
      )
      return lower_one_hot + upper_one_hot
  
file_path: rlax/rlax/_src/transforms.py
imports: import chex
  import jax
  import jax.numpy
  from rlax._src import base
label_desc: Use image node_id 13 for calling a global function with example usage: transform_to_2hot_nonlinear(scalar, bins) and returns: unknown