code: def to_dimension_numbers(
      num_spatial_dims: int,
      channels_last: bool,
      transpose: bool,
  ) -> lax.ConvDimensionNumbers:
      """Create a `lax.ConvDimensionNumbers` for the given inputs."""
      num_dims = num_spatial_dims + 2
  
      if channels_last:
          spatial_dims = tuple(range(1, num_dims - 1))
          image_dn = (0, num_dims - 1) + spatial_dims
      else:
          spatial_dims = tuple(range(2, num_dims))
          image_dn = (0, 1) + spatial_dims
  
      if transpose:
          kernel_dn = (num_dims - 2, num_dims - 1) + tuple(
              range(num_dims - 2)
          )
      else:
          kernel_dn = (num_dims - 1, num_dims - 2) + tuple(
              range(num_dims - 2)
          )
  
      return lax.ConvDimensionNumbers(
          lhs_spec=image_dn, rhs_spec=kernel_dn, out_spec=image_dn
      )
  
imports: from collections.abc import Sequence
  from typing import Optional, Union
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import pad
  from haiku._src import utils
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: to_dimension_numbers(num_spatial_dims, channels_last, transpose) and returns: lax