code: def _fftconvolve_unbatched(
      in1: Array, in2: Array, mode: str
  ) -> Array:
      full_shape = tuple(
          s1 + s2 - 1 for s1, s2 in zip(in1.shape, in2.shape)
      )
      fft_shape = tuple(osp_fft_next_fast_len(s) for s in full_shape)
  
      if mode == "valid":
          no_swap = all(
              s1 >= s2 for s1, s2 in zip(in1.shape, in2.shape)
          )
          swap = all(s1 <= s2 for s1, s2 in zip(in1.shape, in2.shape))
          if not (no_swap or swap):
              raise ValueError(
                  "For 'valid' mode, One input must be at least as "
                  "large as the other in every dimension."
              )
          if swap:
              in1, in2 = in2, in1
  
      if jnp.iscomplexobj(in1):
          fft, ifft = jnp.fft.fftn, jnp.fft.ifftn
      else:
          fft, ifft = jnp.fft.rfftn, jnp.fft.irfftn
      sp1 = fft(in1, fft_shape)
      sp2 = fft(in2, fft_shape)
      conv = ifft(sp1 * sp2, fft_shape)
  
      if mode == "full":
          out_shape = full_shape
      elif mode == "same":
          out_shape = in1.shape
      elif mode == "valid":
          out_shape = tuple(
              s1 - s2 + 1 for s1, s2 in zip(in1.shape, in2.shape)
          )
      else:
          raise ValueError(f"Unrecognized {mode=}")
  
      start_indices = tuple(
          (full_size - out_size) // 2
          for full_size, out_size in zip(full_shape, out_shape)
      )
      return lax.dynamic_slice(conv, start_indices, out_shape)
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  from functools import partial
  import math
  import operator
  from typing import Callable
  import warnings
  import numpy
  import scipy.signal
  from scipy.fft import next_fast_len
  import jax
  import jax.numpy.fft
  import jax.numpy
  from jax import lax
  from jax._src.api_util import _ensure_index_tuple
  from jax._src import dtypes
  from jax._src.lax.lax import PrecisionLike
  from jax._src.numpy import linalg
  from jax._src.numpy.util import check_arraylike, _wraps, promote_dtypes_inexact, promote_dtypes_complex
  from jax._src.third_party.scipy import signal_helper
  from jax._src.typing import Array, ArrayLike
  from jax._src.util import canonicalize_axis, tuple_delete, tuple_insert
label_desc: Use image node_id 2 for calling a global function with example usage: _fftconvolve_unbatched(in1, in2, mode) and returns: lax