code: def weighted_mean_squared_error(
      predictions: jnp.ndarray,
      targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
      axis: Optional[Union[int, Tuple[int, ...]]] = None,
  ) -> jnp.ndarray:
      """Weighted mean of weighted_squared_error.
  
      Args:
        predictions: Output of model in shape [batch, ..., num_features].
        targets: Targets of shape [batch, ..., num_features].
        weights:  None or array of shape [batch,] This is the weight to apply to the
          loss  computed for each example in the batch. Can be used to ignore padded
          examples in the batch.
        axis: The axis (or axes) to compute the loss over. If not specified, all
          dimensions besides the leading batch dimension are used.
  
      Returns:
        The averaged mean squared error of all the examples in the given batch as a
        scalar.
      """
      unnormalized_mse = weighted_squared_error(
          predictions=predictions,
          targets=targets,
          weights=weights,
          axis=axis,
      )
  
      if weights is not None:
          # Divide by sum of the broadcasted weights:
          broadcasted_shape = weights.shape + (1,) * (
              unnormalized_mse.ndim - weights.ndim
          )
          broadcasted_weights = jax.lax.broadcast_in_dim(
              weights,
              shape=broadcasted_shape,
              broadcast_dimensions=tuple(range(weights.ndim)),
          )
          normalization = jnp.sum(
              broadcasted_weights * jnp.ones(unnormalized_mse.shape)
          )
      else:
          # Divide by number of examples:
          normalization = unnormalized_mse.size
      return jnp.sum(unnormalized_mse) / (normalization + 1e-8)
  
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 18 for calling a global function with example usage: weighted_mean_squared_error(predictions, targets, weights, axis) and returns: unknown