code: def _conv_dimension_numbers(input_shape):
      """Computes the dimension numbers based on the input shape."""
      ndim = len(input_shape)
      lhs_spec = (0, ndim - 1) + tuple(range(1, ndim - 1))
      rhs_spec = (ndim - 1, ndim - 2) + tuple(range(0, ndim - 2))
      out_spec = lhs_spec
      return lax.ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)
  
imports: import dataclasses
  from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Union
  import jax
  import jax.numpy
  import numpy
  from jax import eval_shape, lax
  from jax.core import ShapedArray
  from flax.core import meta
  from flax.linen import initializers
  from flax.linen.dtypes import promote_dtype
  from flax.linen.module import Module, compact
label_desc: Use image node_id 3 for calling a global function with example usage: _conv_dimension_numbers(input_shape) and returns: lax