code: def likelihood(predictions: Array, targets: Array) -> Array:
      """Calculates the likelihood of predictions wrt targets.
  
      Args:
        predictions: a vector of arbitrary shape.
        targets: a vector of shape compatible with predictions.
  
      Returns:
        a vector of same shape of `predictions`.
      """
      chex.assert_type([predictions, targets], float)
      likelihood_vals = predictions**targets * (
          1.0 - predictions
      ) ** (1.0 - targets)
      # Note: 0**0 evaluates to NaN on TPUs, manually set these cases to 1.
      filter_indices = jnp.logical_or(
          jnp.logical_and(targets == 1, predictions == 1),
          jnp.logical_and(targets == 0, predictions == 0),
      )
      return jnp.where(filter_indices, 1, likelihood_vals)
  
file_path: rlax/rlax/_src/losses.py
imports: import functools
  from typing import Optional, Union
  import chex
  import jax
  import jax.numpy
  from rlax._src import general_value_functions
  from rlax._src import value_learning
label_desc: Use image node_id 2 for calling a global function with example usage: likelihood(predictions, targets) and returns: jnp