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)
  
file_path: flax/flax/core/nn/linear.py
imports: from collections.abc import Iterable
  import jax.numpy
  import numpy
  from jax import lax
  from flax import struct
  from flax.core import Scope
  from flax.linen import initializers
label_desc: Use image node_id 4 for calling a global function with example usage: _conv_dimension_numbers(input_shape) and returns: lax