code: def weighted_squared_error(
      predictions: jnp.ndarray,
      targets: jnp.ndarray,
      weights: Optional[jnp.ndarray] = None,
      axis: Optional[Union[int, Tuple[int, ...]]] = None,
  ) -> jnp.ndarray:
      """Computes weighted squared error given predictions and targets.
  
      This computes the squared_error of examples in a single, potentially
      padded minibatch. If the minibatch is padded (that is it contains null
      examples) it is assumed that weights is a binary mask where 0 indicates that
      the example is null.
  
      Args:
        predictions: Output of model in shape shape [batch, ..., n_features].
        targets: Array of shape [batch, ..., n_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 mean squared error for each example in the given batch. The output shape
        is [batch,].
      """
      if predictions.ndim != targets.ndim:
          raise ValueError(
              "Incorrect shapes. Got shape %s predictions and %s targets"
              % (str(predictions.shape), str(targets.shape))
          )
      if axis is None:
          # Sum over all features in each example in the batch:
          axis = tuple(range(1, predictions.ndim))
  
      error = targets - predictions
      loss = jnp.square(error)
      loss = jnp.sum(loss, axis=axis)
      if weights is not None:
          loss = apply_weights(loss, weights)
      return loss
  
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 17 for calling a global function with example usage: weighted_squared_error(predictions, targets, weights, axis) and returns: loss