code: def combine_masks(
      *masks: Optional[Array], dtype: Dtype = jnp.float32
  ) -> Array:
      """Combine attention masks.
  
      Args:
        *masks: set of attention mask arguments to combine, some can be None.
        dtype: dtype for the returned mask.
  
      Returns:
        Combined mask, reduced by logical and, returns None if no masks given.
      """
      masks_list = [m for m in masks if m is not None]
      if not masks_list:
          return None
      assert all(
          map(lambda x: x.ndim == masks_list[0].ndim, masks_list)
      ), f"masks must have same rank: {tuple(map(lambda x: x.ndim, masks_list))}"
      mask, *other_masks = masks_list
      for other_mask in other_masks:
          mask = jnp.logical_and(mask, other_mask)
      return mask.astype(dtype)
  
imports: from __future__ import annotations
  import functools
  from typing import Any, Callable, Optional, Tuple, overload
  import jax
  import jax.numpy
  from jax import lax, random
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import rnglib
  from flax.experimental.nnx.nnx import flaglib
  from flax.experimental.nnx.nnx.module import Module, first_from
  from flax.experimental.nnx.nnx.nn import initializers
  from flax.experimental.nnx.nnx.nn.dtypes import promote_dtype
  from flax.experimental.nnx.nnx.nn.linear import DotGeneralT, LinearGeneral, PrecisionLike, default_kernel_init
  from flax.experimental.nnx.nnx.nn.normalization import LayerNorm
label_desc: Use image node_id 5 for calling a global function with example usage: combine_masks() and returns: mask, None