code: def apply_weights(
      output: jnp.ndarray, weights: jnp.ndarray
  ) -> jnp.ndarray:
      """Applies given weights of the inputs in the minibatch to outputs.
  
      Note that weights can be per example (i.e. of shape `[batch,]`) or per
      pixel/token (i.e. of shape `[batch, height, width]` or
      `[batch, len]`) so we need to broadcast it to the output shape.
  
      Args:
        output: Computed output, which can be loss or the correctly classified
          examples, etc.
        weights: Weights of inputs in the batch, which can be None or array of shape
          [batch, ...].
  
      Returns:
        Weighted output.
      """
      if output.ndim < weights.ndim:
          raise ValueError(
              "Output rank should be higher or equal to weights rank."
          )
      desired_weights_shape = weights.shape + (1,) * (
          output.ndim - weights.ndim
      )
      weights = jax.lax.broadcast_in_dim(
          weights,
          shape=desired_weights_shape,
          broadcast_dimensions=tuple(range(weights.ndim)),
      )
      # Scale the outputs with weights.
      return output * weights
  
imports: import functools
  from typing import Optional, Any, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: apply_weights(output, weights) and returns: unknown