code: def _conv_output_shape(
      spatial_shape: jnp.ndarray,
      kernel_size: Tuple[int, ...],
      input_dilation: Optional[Tuple[int, ...]],
      kernel_dilation: Optional[Tuple[int, ...]],
      strides: Optional[Tuple[int, ...]],
      padding: Union[str, Sequence[Tuple[int, int]]],
  ) -> jnp.ndarray:
      """Convenience wrapper function for inferring the convolution output shape.
  
      Args:
        spatial_shape: Input (lhs) shapes for which the output shapes should be
          inferred as array with dimensions (batch, spatial_dims, dims).
        kernel_size: Covolution kernel size (i.e. rhs shape).
        input_dilation: Input (lhs) dilation.
        kernel_dilation: Convolution kernel (rhs) dilation.
        strides: Convolution (rhs) stride.
        padding: Input (lhs) padding.
  
      Returns:
        Inferred convolution output shapes as array with dimensions
        (batch, len(spatial_dims)).
      """
      strides = strides or (1,) * len(kernel_size)
      if input_dilation is not None:
          spatial_shape = _dilate_shape(spatial_shape, input_dilation)
      if kernel_dilation is not None:
          kernel_size = tuple(
              (k - 1) * r + 1
              for k, r in zip(kernel_size, kernel_dilation)
          )
      spatial_shape = jnp.concatenate(
          [
              jnp.ones((spatial_shape.shape[0], 2), dtype=jnp.int32),
              spatial_shape,
          ],
          axis=-1,
      )
      out_shape = conv_shape_tuple(
          lhs_shape=spatial_shape,
          rhs_shape=(1, 1) + kernel_size,
          strides=strides,
          pads=padding,
      )
      return out_shape[:, 2:]
  
imports: import functools
  from typing import Optional, Tuple, Union, Callable, Any, Sequence, List, Iterable
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 8 for calling a global function with example usage: _conv_output_shape(spatial_shape, kernel_size, input_dilation, kernel_dilation, strides, padding) and returns: out_shape