code: def apply_spatial_mask(
      inputs: jnp.ndarray,
      spatial_shape: jnp.ndarray,
      value: float = 0.0,
  ) -> jnp.ndarray:
      """Construct and apply spatial mask to the inputs.
  
      Args:
        inputs: Input tensor with dimensions [batch, spatial_dim, dim] to which the
          mask should be applied.
        spatial_shape: Per-input spatial shape of *unpadded* input data with
          dimensions (batch, len(spatial_dims)).
        value: Value to use for the mask (default: 0).
  
      Returns:
        Inputs with masking spatial masking applied to them.
      """
      assert inputs.shape[0] == spatial_shape.shape[0]
      assert spatial_shape.shape[1] == inputs.ndim - 2
      mask = mask_from_spatial(
          inputs.shape[1:-1], spatial_shape, per_axis=False
      )
      mask = jnp.expand_dims(mask, axis=-1)
      inputs = jnp.where(mask, inputs, value)
      return inputs
  
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 5 for calling a global function with example usage: apply_spatial_mask(inputs, spatial_shape, value) and returns: inputs