code: def _infer_shape(
      x: jax.Array,
      size: Union[int, Sequence[int]],
      channel_axis: Optional[int] = -1,
  ) -> tuple[int, ...]:
      """Infer shape for pooling window or strides."""
      if isinstance(size, int):
          if channel_axis and not 0 <= abs(channel_axis) < x.ndim:
              raise ValueError(
                  f"Invalid channel axis {channel_axis} for {x.shape}"
              )
          if channel_axis and channel_axis < 0:
              channel_axis = x.ndim + channel_axis
          return (1,) + tuple(
              size if d != channel_axis else 1 for d in range(1, x.ndim)
          )
      elif len(size) < x.ndim:
          # Assume additional dimensions are batch dimensions.
          return (1,) * (x.ndim - len(size)) + tuple(size)
      else:
          assert x.ndim == len(size)
          return tuple(size)
  
imports: from collections.abc import Sequence
  from typing import Optional, Union
  import warnings
  from haiku._src import module
  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: _infer_shape(x, size, channel_axis) and returns: unknown, unknown, tuple