code: def entropy_loss(logits, *args):
      """rlax.entropy_loss, but with sum(loss) and [T, B, ...] inputs."""
      mean_per_batch = jax.vmap(rlax.entropy_loss, in_axes=1)(
          logits, *args
      )
      total_loss_per_batch = mean_per_batch * logits.shape[0]
      return jnp.sum(total_loss_per_batch)
  
imports: import functools
  import itertools
  import queue
  import threading
  import warnings
  import dm_env
  import haiku
  from examples.impala import agent
  from examples.impala import util
  import jax
  from jax.example_libraries import optimizers
  import jax.numpy
  import numpy
  import optax
  import rlax
label_desc: Use image node_id 2 for calling a global function with example usage: entropy_loss(logits) and returns: jnp