code: def _weighted_average(
      x: jnp.ndarray, y: jnp.ndarray, lambdas: jnp.ndarray
  ) -> jnp.ndarray:
      return lambdas * x + (1.0 - lambdas) * y
  
file_path: acme/acme/agents/jax/ail/losses.py
imports: import functools
  from typing import Callable, Dict, Optional, Tuple
  from acme import types
  from acme.jax import networks
  import jax
  import jax.numpy
  import tensorflow_probability
  import tree
label_desc: Use image node_id 2 for calling a global function with example usage: _weighted_average(x, y, lambdas) and returns: unknown