code: def unnormalize_linear(
      state: PopArtState, inputs: Array, indices: Array
  ) -> PopArtOutput:
      """Selects and unnormalizes output of a Linear.
  
      Args:
        state: The PopArt summary stats.
        inputs: The (normalized) output of the Linear that we apply PopArt to.
        indices: Which indices of `inputs` to use.
  
      Returns:
        PopArtOutput, a tuple of the normalized and unnormalized PopArt values.
      """
      assert jnp.issubdtype(indices.dtype, jnp.integer)
      assert indices.shape == inputs.shape[:-1]
      normalized = jnp.take_along_axis(
          inputs, indices[..., None], axis=-1
      )
      normalized = jnp.squeeze(normalized, axis=-1)
      return PopArtOutput(
          normalized, unnormalize(state, normalized, indices)
      )
  
file_path: rlax/rlax/_src/pop_art.py
imports: import collections
  from typing import Mapping, Tuple
  import chex
  import jax
  import jax.numpy
  from rlax._src import base
label_desc: Use image node_id 4 for calling a global function with example usage: unnormalize_linear(state, inputs, indices) and returns: PopArtOutput