code: def mask_from_spatial(
      padded_shape: Tuple[int, ...],
      spatial_shape: jnp.ndarray,
      per_axis: bool = False,
  ) -> Union[jnp.ndarray, List[jnp.ndarray]]:
      """Create a spatial mask for a given padded shape and spatial size.
  
      Args:
        padded_shape: Shape of the spatial dimensions padded data that needs to be
          masked.
        spatial_shape: Per-element unpadded spatial size of the data with dimensions
          (batch, len(padded_shape)).
        per_axis: If True, per list of per-spatial-dim masks is returned instead of
          a single mask; that should enable some memory savings as it allows for
          shape boradcasting to happen only when it is required.
  
      Returns:
        The spatial mask.
      """
      assert spatial_shape.ndim == 2
      assert len(padded_shape) == spatial_shape.shape[1]
      ndim = spatial_shape.shape[1]
  
      masks = []
      for i in range(ndim):
          # Construct per-axis mask and then broadcast if asked for.
          mask = jnp.arange(0, padded_shape[i], dtype=jnp.int32)
          mask = jnp.reshape(
              mask,
              (1,) * (i + 1)
              + (padded_shape[i],)
              + (1,) * (ndim - i - 1),
          )
          threshold = spatial_shape[:, i]
          threshold = jnp.reshape(
              threshold, (spatial_shape.shape[0],) + (1,) * ndim
          )
          mask = mask < threshold
          masks.append(mask)
  
      if not per_axis:
          masks = functools.reduce(jnp.logical_and, masks)
      return masks
  
imports: import functools
  from typing import Optional, Tuple, Union, Callable, Any, Sequence, List, Iterable
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 6 for calling a global function with example usage: mask_from_spatial(padded_shape, spatial_shape, per_axis) and returns: masks