code: def nce_loss(
      scores: jnp.ndarray,  # Shape: (N, N)
      where: Optional[
          jnp.ndarray
      ] = None,  # Shape: broadcastable with (N,)
      initial: Optional[float] = None,
  ) -> jnp.ndarray:  # Shape: (N,)
      return _rows_to_columns_nce_loss(
          scores, where=where, initial=initial
      ) + _rows_to_columns_nce_loss(
          scores.T, where=where, initial=initial
      )
  
imports: from typing import Optional
  from jax import nn
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: nce_loss(scores, where, initial) and returns: unknown