code: def transform_from_2hot_nonlinear(probs: Array, bins: Array) -> Array:
      """Transforms from a categorical distribution to a scalar."""
      scalar = jnp.sum(probs * jnp.expand_dims(bins, 0), -1)
      return scalar
  
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 14 for calling a global function with example usage: transform_from_2hot_nonlinear(probs, bins) and returns: scalar