code: def conv_shape_tuple(
      lhs_shape: jnp.ndarray,
      rhs_shape: Tuple[int, ...],
      strides: Tuple[int, ...],
      pads: Union[str, Sequence[Tuple[int, int]]],
  ) -> jnp.ndarray:
      """Compute the shape of a conv given input shapes in canonical order.
  
      Args:
        lhs_shape: Input (lhs) shapes for which the output shapes should be inferred
          as array with dimensions (batch, spatial_dims, dims).
        rhs_shape: Covolution kernel size (i.e. rhs shape).
        strides: Convolution (rhs) stride.
        pads: Input (lhs) padding.
  
      Returns:
        Inferred convolution output shapes as array with dimensions
        (batch, len(spatial_dims)).
      """
      if isinstance(pads, str):
          pads = padtype_to_pads(
              lhs_shape[:, 2:], rhs_shape[2:], strides, pads
          )
      else:
          pads = jnp.expand_dims(jnp.array(pads), axis=0)
  
      if pads.shape[1] != lhs_shape.shape[1] - 2:
          msg = "Wrong number of explicit pads for convolution: expected {}, got {}."
          raise TypeError(
              msg.format(lhs_shape.shape[1] - 2, pads.shape[1])
          )
      lhs_padded = jnp.add(lhs_shape[:, 2:], jnp.sum(pads, axis=2))
  
      rhs_shape = jnp.array(rhs_shape, dtype=jnp.int32)
      strides = jnp.array(strides, dtype=jnp.int32)
      out_space = (
          jnp.floor_divide(
              jnp.subtract(lhs_padded, rhs_shape[2:]), strides
          )
          + 1
      )
      out_space = jnp.maximum(0, out_space)
      out_shape = jnp.stack(
          [
              lhs_shape[:, 0],
              jnp.full((lhs_shape.shape[0],), rhs_shape[0]),
          ],
          axis=-1,
      )
      return jnp.concatenate([out_shape, out_space], axis=-1)
  
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 11 for calling a global function with example usage: conv_shape_tuple(lhs_shape, rhs_shape, strides, pads) and returns: jnp