code: def _validate_spatial_dimensions(
      lhs: TfVal,
      lhs_shape: core.Shape,
      rhs: TfVal,
      rhs_shape: core.Shape,
  ):
      """Check spatial dimension support."""
      jax2tf._assert_matching_abstract_shape(lhs, lhs_shape)
      jax2tf._assert_matching_abstract_shape(rhs, rhs_shape)
  
      nr_spatial_dimensions = len(lhs_shape) - 2
      # Currently we only support 1D+2D convolutions because it keeps the code
      # relatively simple and covers most cases.
      if nr_spatial_dimensions > 2:
          raise _conv_error(
              "We only support 1D or 2D convolutions, but found "
              f"{nr_spatial_dimensions}."
          )
  
imports: from __future__ import annotations
  import builtins
  from collections.abc import Sequence
  import dataclasses
  from functools import partial, wraps
  import math
  import string
  from typing import Any, Callable, Optional
  from jax._src import core
  from jax import lax
  from jax._src.lax import slicing
  from jax._src import dtypes
  from jax._src import util
  from jax.experimental.jax2tf import jax2tf
  import numpy
  import tensorflow
label_desc: Use image node_id 10 for calling a global function with example usage: _validate_spatial_dimensions(lhs, lhs_shape, rhs, rhs_shape) without return types