Pooling#

class serket.nn.MaxPool1D(kernel_size, strides=1, *, padding='valid')[source]#

1D Max Pooling layer

Parameters:
  • kernel_size (Union[int, Sequence[int]]) – size of the kernel

  • strides (Union[int, Sequence[int]]) – strides of the kernel

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – padding of the kernel (valid, same) or tuple of ints

Example

>>> import jax
>>> import jax.numpy as jnp
>>> import serket as sk
>>> layer = sk.nn.MaxPool1D(kernel_size=2, strides=2)
>>> x = jnp.arange(1, 11).reshape(1, 10).astype(jnp.float32)
>>> print(layer(x))
[[ 2.  4.  6.  8. 10.]]
__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.MaxPool2D(kernel_size, strides=1, *, padding='valid')[source]#

2D Max Pooling layer

Parameters:
  • kernel_size (Union[int, Sequence[int]]) – size of the kernel

  • strides (Union[int, Sequence[int]]) – strides of the kernel

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – padding of the kernel (valid, same) or tuple of ints

Example

>>> import jax
>>> import jax.numpy as jnp
>>> import serket as sk
>>> layer = sk.nn.MaxPool2D(kernel_size=2, strides=2)
>>> x = jnp.arange(1, 17).reshape(1, 4, 4).astype(jnp.float32)
>>> print(layer(x))
[[[ 6.  8.]
  [14. 16.]]]
__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.MaxPool3D(kernel_size, strides=1, *, padding='valid')[source]#

3D Max Pooling layer

Parameters:
  • kernel_size (Union[int, Sequence[int]]) – size of the kernel

  • strides (Union[int, Sequence[int]]) – strides of the kernel

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – padding of the kernel (valid, same) or tuple of ints

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.AvgPool1D(kernel_size, strides=1, *, padding='valid')[source]#

1D Average Pooling layer

Parameters:
  • kernel_size (Union[int, Sequence[int]]) – size of the kernel

  • strides (Union[int, Sequence[int]]) – strides of the kernel

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – padding of the kernel (valid, same) or tuple of ints

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.AvgPool2D(kernel_size, strides=1, *, padding='valid')[source]#

2D Average Pooling layer

Parameters:
  • kernel_size (Union[int, Sequence[int]]) – size of the kernel

  • strides (Union[int, Sequence[int]]) – strides of the kernel

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – padding of the kernel (valid, same) or tuple of ints

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.AvgPool3D(kernel_size, strides=1, *, padding='valid')[source]#

3D Average Pooling layer

Parameters:
  • kernel_size (Union[int, Sequence[int]]) – size of the kernel

  • strides (Union[int, Sequence[int]]) – strides of the kernel

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – padding of the kernel (valid, same) or tuple of ints

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.GlobalAvgPool1D(keepdims=True)[source]#

1D Global Average Pooling layer

Parameters:

keepdims (bool) – whether to keep the dimensions or not

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.GlobalAvgPool2D(keepdims=True)[source]#

2D Global Average Pooling layer

Parameters:

keepdims (bool) – whether to keep the dimensions or not

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.GlobalAvgPool3D(keepdims=True)[source]#

3D Global Average Pooling layer

Parameters:

keepdims (bool) – whether to keep the dimensions or not

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.GlobalMaxPool1D(keepdims=True)[source]#

1D Global Max Pooling layer

Parameters:

keepdims (bool) – whether to keep the dimensions or not

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.GlobalMaxPool2D(keepdims=True)[source]#

2D Global Max Pooling layer

Parameters:

keepdims (bool) – whether to keep the dimensions or not

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.GlobalMaxPool3D(keepdims=True)[source]#

3D Global Max Pooling layer

Parameters:

keepdims (bool) – whether to keep the dimensions or not

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.LPPool1D(norm_type, kernel_size, strides=1, *, padding='valid')[source]#

1D Lp pooling to the input.

Parameters:
  • norm_type (float) – norm type

  • kernel_size (Union[int, Sequence[int]]) – size of the kernel

  • strides (Union[int, Sequence[int]]) – strides of the kernel

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – padding of the kernel

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.LPPool2D(norm_type, kernel_size, strides=1, *, padding='valid')[source]#

2D Lp pooling to the input.

Parameters:
  • norm_type (float) – norm type

  • kernel_size (Union[int, Sequence[int]]) – size of the kernel

  • strides (Union[int, Sequence[int]]) – strides of the kernel

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – padding of the kernel

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.LPPool3D(norm_type, kernel_size, strides=1, *, padding='valid')[source]#

3D Lp pooling to the input.

Parameters:
  • norm_type (float) – norm type

  • kernel_size (Union[int, Sequence[int]]) – size of the kernel

  • strides (Union[int, Sequence[int]]) – strides of the kernel

  • padding (Union[str, int, Sequence[int], Sequence[Tuple[int, int]]]) – padding of the kernel

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.AdaptiveAvgPool1D(output_size)[source]#

1D Adaptive Average Pooling layer

Parameters:

output_size (tuple[int, ...]) – size of the output

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.AdaptiveAvgPool2D(output_size)[source]#

2D Adaptive Average Pooling layer

Parameters:

output_size (tuple[int, ...]) – size of the output

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.AdaptiveAvgPool3D(output_size)[source]#

3D Adaptive Average Pooling layer

Parameters:

output_size (tuple[int, ...]) – size of the output

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.AdaptiveMaxPool1D(output_size)[source]#

1D Adaptive Max Pooling layer

Parameters:

output_size (tuple[int, ...]) – size of the output

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.AdaptiveMaxPool2D(output_size)[source]#

2D Adaptive Max Pooling layer

Parameters:

output_size (tuple[int, ...]) – size of the output

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.AdaptiveMaxPool3D(output_size)[source]#

3D Adaptive Max Pooling layer

Parameters:

output_size (tuple[int, ...]) – size of the output

__call__(input)#

Call self as a function.

Parameters:

input (Array)

Return type:

Array

serket.nn.adaptive_avg_pool_nd(input, out_dim)[source]#

Adaptive average pooling operation

Parameters:
  • input (Array) – channeled input of shape (channels, spatial_dims)

  • out_dim (Sequence[int]) – output dimension. accepts a sequence of ints for each spatial dimension

Return type:

Array

Example

>>> import jax
>>> import jax.numpy as jnp
>>> import serket as sk
>>> input = jnp.ones((2, 25, 25))
>>> out_dim = (13, 13)
>>> output = sk.nn.adaptive_avg_pool_nd(input, out_dim)
>>> print(output.shape)
(2, 13, 13)
serket.nn.adaptive_max_pool_nd(input, out_dim)[source]#

Adaptive max pooling operation

Parameters:
  • input (Array) – channeled input of shape (channels, spatial_dims)

  • out_dim (Sequence[int]) – output dimension. accepts a sequence of ints for each spatial dimension

Return type:

Array

Example

>>> import jax
>>> import jax.numpy as jnp
>>> import serket as sk
>>> input = jnp.ones((2, 25, 25))
>>> out_dim = (13, 13)
>>> output = sk.nn.adaptive_max_pool_nd(input, out_dim)
>>> print(output.shape)
(2, 13, 13)
serket.nn.avg_pool_nd(input, kernel_size, strides, padding)[source]#

Average pooling operation

Parameters:
  • input (Array) – channeled input of shape (channels, spatial_dims)

  • kernel_size (Sequence[int]) – size of the kernel. accepts tuple of ints for each spatial dimension

  • strides (Sequence[int]) – strides of the kernel. accepts tuple of ints for each spatial dimension

  • padding (Sequence[tuple[int, int]]) – padding of the kernel. accepts tuple of tuples of two ints for each spatial dimension for each side of the input

Return type:

Array

Example

>>> import jax
>>> import jax.numpy as jnp
>>> import serket as sk
>>> kernel_size = (3, 3)
>>> strides = (2, 2)
>>> input = jnp.ones((2, 25, 25))
>>> padding = ((1, 1), (1, 1)) # pad 1 on each side of the spatial dimensions
>>> output = sk.nn.avg_pool_nd(input, kernel_size, strides, padding)
>>> print(output.shape)
(2, 13, 13)
serket.nn.lp_pool_nd(input, norm_type, kernel_size, strides, padding)[source]#

Lp pooling operation

Parameters:
  • input (Array) – channeled input of shape (channels, spatial_dims)

  • norm_type (float) – norm type as a float

  • kernel_size (Sequence[int]) – size of the kernel. accepts tuple of ints for each spatial dimension

  • strides (Sequence[int]) – strides of the kernel. accepts tuple of ints for each spatial dimension

  • padding (Sequence[tuple[int, int]]) – padding of the kernel. accepts tuple of tuples of two ints for each spatial dimension for each side of the input

Return type:

Array

Example

>>> import jax
>>> import jax.numpy as jnp
>>> import serket as sk
>>> kernel_size = (3, 3)
>>> strides = (2, 2)
>>> input = jnp.ones((2, 25, 25))
>>> norm_type = 2
>>> padding = ((1, 1), (1, 1)) # pad 1 on each side of the spatial dimensions
>>> output = sk.nn.lp_pool_nd(input, norm_type, kernel_size, strides, padding)
>>> print(output.shape)
(2, 13, 13)
serket.nn.max_pool_nd(input, kernel_size, strides, padding)[source]#

Max pooling operation

Parameters:
  • input (Array) – channeled input of shape (channels, spatial_dims)

  • kernel_size (Sequence[int]) – size of the kernel. accepts a sequence of ints for each spatial dimension

  • strides (Sequence[int]) – strides of the kernel. accepts a sequence of ints for each spatial dimension

  • padding (Sequence[tuple[int, int]]) – padding of the kernel. accepts a sequence of tuples of two ints for each spatial dimension for each side of the input

Return type:

Array

Example

>>> import jax
>>> import jax.numpy as jnp
>>> import serket as sk
>>> kernel_size = (3, 3)
>>> strides = (2, 2)
>>> input = jnp.ones((2, 25, 25))
>>> padding = ((1, 1), (1, 1)) # pad 1 on each side of the spatial dimensions
>>> output = sk.nn.max_pool_nd(input, kernel_size, strides, padding)
>>> print(output.shape)
(2, 13, 13)