class_docstr: General N-dimensional downsampling over a sliding window.
class_docstr_tok: ["General","N-dimensional","downsampling","over","a","sliding","window","."]
docstr: **Arguments:**
  
  - `init`: The initial value for the reduction.
  - `operation`: The operation applied to the inputs of each window.
  - `num_spatial_dims`: The number of spatial dimensions.
  - `kernel_size`: The size of the convolutional kernel.
  - `stride`: The stride of the convolution.
  - `padding`: The amount of padding to apply before and after each
      spatial dimension.
  - `use_ceil`: If `True`, then `ceil` is used to compute the final output
      shape instead of `floor`. For `ceil`, if required, extra padding is added.
      Defaults to `False`.
  
  !!! info
  
      In order for `Pool` to be differentiable, `operation(init, x) == x` needs to
      be true for all finite `x`. For further details see
      [https://www.tensorflow.org/xla/operation_semantics#reducewindow](https://www.tensorflow.org/xla/operation_semantics#reducewindow)
      and [https://github.com/google/jax/issues/7718](https://github.com/google/jax/issues/7718).
code: def __init__(
      self,
      init: Union[int, float, Array],
      operation: Callable[[Array, Array], Array],
      num_spatial_dims: int,
      kernel_size: Union[int, Sequence[int]],
      stride: Union[int, Sequence[int]] = 1,
      padding: Union[int, Sequence[int], Sequence[tuple[int, int]]] = 0,
      use_ceil: bool = False,
  ):
      """**Arguments:**
  
      - `init`: The initial value for the reduction.
      - `operation`: The operation applied to the inputs of each window.
      - `num_spatial_dims`: The number of spatial dimensions.
      - `kernel_size`: The size of the convolutional kernel.
      - `stride`: The stride of the convolution.
      - `padding`: The amount of padding to apply before and after each
          spatial dimension.
      - `use_ceil`: If `True`, then `ceil` is used to compute the final output
          shape instead of `floor`. For `ceil`, if required, extra padding is added.
          Defaults to `False`.
  
      !!! info
  
          In order for `Pool` to be differentiable, `operation(init, x) == x` needs to
          be true for all finite `x`. For further details see
          [https://www.tensorflow.org/xla/operation_semantics#reducewindow](https://www.tensorflow.org/xla/operation_semantics#reducewindow)
          and [https://github.com/google/jax/issues/7718](https://github.com/google/jax/issues/7718).
      """  # noqa: E501
  
      self.operation = operation
      self.init = init
      self.num_spatial_dims = num_spatial_dims
      self.use_ceil = use_ceil
  
      if isinstance(kernel_size, int):
          self.kernel_size = (kernel_size,) * num_spatial_dims
      elif isinstance(kernel_size, Sequence):
          self.kernel_size = tuple(kernel_size)
      else:
          raise ValueError(
              "`kernel_size` must either be an int or tuple of length "
              f"{num_spatial_dims} containing ints."
          )
  
      if isinstance(stride, int):
          self.stride = (stride,) * num_spatial_dims
      elif isinstance(stride, Sequence):
          self.stride = tuple(stride)
      else:
          raise ValueError(
              "`stride` must either be an int or tuple of length "
              f"{num_spatial_dims} containing ints."
          )
  
      if isinstance(padding, int):
          self.padding = tuple(
              (padding, padding) for _ in range(num_spatial_dims)
          )
      elif (
          isinstance(padding, Sequence)
          and len(padding) == num_spatial_dims
      ):
          if all_sequences(padding):
              self.padding = tuple(padding)
          else:
              self.padding = tuple((p, p) for p in padding)
      else:
          raise ValueError(
              "`padding` must either be an int or tuple of length "
              f"{num_spatial_dims} containing ints or tuples of length 2."
          )
  
imports: import math
  from collections.abc import Callable, Sequence
  from typing import Optional, Union
  import jax
  import jax.lax
  import jax.numpy
  import jax.random
  from jaxtyping import Array, PRNGKeyArray
  from .._module import field, Module
  from ._misc import all_sequences
label_desc: Use image node_id 1 to create a new Pool object from inherited base classes: Module with example: obj = Pool(init, operation, num_spatial_dims, kernel_size, stride, padding, use_ceil)