Reshaping

class serket.nn.CenterCrop1D(size)

Crops a 1D input to the given size at the center.

Parameters:

size (int | tuple[int, ...]) – The size of the output image. accepts a single int.

Example

>>> import serket as sk
>>> import jax.numpy as jnp
>>> layer = sk.nn.CenterCrop1D(4)
>>> input = jnp.arange(1, 13).reshape(1, 12)
>>> print(input)
[[ 1  2  3  4  5  6  7  8  9 10 11 12]]
>>> print(layer(input))
[[5 6 7 8]]

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.CenterCrop2D(size)

Crop the center of a channel-first image.

Parameters:

size (int | tuple[int, ...]) – The size of the output image. accepts a single int or a tuple of two ints.

Example

>>> import serket as sk
>>> import jax.numpy as jnp
>>> layer = sk.nn.CenterCrop2D(4)
>>> input = jnp.arange(1, 145).reshape(1, 12, 12)
>>> print(input)
[[[  1   2   3   4   5   6   7   8   9  10  11  12]
  [ 13  14  15  16  17  18  19  20  21  22  23  24]
  [ 25  26  27  28  29  30  31  32  33  34  35  36]
  [ 37  38  39  40  41  42  43  44  45  46  47  48]
  [ 49  50  51  52  53  54  55  56  57  58  59  60]
  [ 61  62  63  64  65  66  67  68  69  70  71  72]
  [ 73  74  75  76  77  78  79  80  81  82  83  84]
  [ 85  86  87  88  89  90  91  92  93  94  95  96]
  [ 97  98  99 100 101 102 103 104 105 106 107 108]
  [109 110 111 112 113 114 115 116 117 118 119 120]
  [121 122 123 124 125 126 127 128 129 130 131 132]
  [133 134 135 136 137 138 139 140 141 142 143 144]]]
>>> print(layer(input))
[[[53 54 55 56]
  [65 66 67 68]
  [77 78 79 80]
  [89 90 91 92]]]

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.CenterCrop3D(size)

Crops a 3D input to the given size at the center.

Parameters:

size (int | tuple[int, ...])

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.RandomCrop1D(size)

Crop a 1D input to the given size at a random start.

Parameters:

size (int | tuple[int, ...]) – size of the slice, either a single int or a tuple of int. accepted values are either a single int or a tuple of int denoting the size.

__call__(input, *, key)

Call self as a function.

Parameters:

• input (Array)

• key (Array)

Return type:

Array

class serket.nn.RandomCrop2D(size)

Crop a 2D input to the given size at a random start.

Parameters:

size (int | tuple[int, ...]) – size of the slice in each axis. accepted values are either a single int or a tuple of two ints denoting the size along each axis.

__call__(input, *, key)

Call self as a function.

Parameters:

• input (Array)

• key (Array)

Return type:

Array

class serket.nn.RandomCrop3D(size)

Crop a 3D input to the given size at a random start.

Parameters:

size (int | tuple[int, ...]) – size of the slice in each axis. accepted values are either a single int or a tuple of three ints denoting the size along each axis.

__call__(input, *, key)

Call self as a function.

Parameters:

• input (Array)

• key (Array)

Return type:

Array

class serket.nn.Upsample1D(scale=1, method='nearest')

Upsample a 1D input to a given size using a given interpolation method.

Parameters:

• scale (int | tuple[int, ...]) – the scale of the output.

• method (Literal['nearest', 'linear', 'cubic', 'lanczos3', 'lanczos5']) –

  Interpolation method defaults to nearest. choices are:

  • nearest: Nearest neighbor interpolation. The values of antialias and precision are ignored.

  • linear, bilinear, trilinear, triangle: Linear interpolation. If antialias is True, uses a triangular filter when downsampling.

  • cubic, bicubic, tricubic: Cubic interpolation, using the keys cubic kernel.

  • lanczos3: Lanczos resampling, using a kernel of radius 3.

  • lanczos5: Lanczos resampling, using a kernel of radius 5.

Example

>>> import serket as sk
>>> import jax.numpy as jnp
>>> layer = sk.nn.Upsample1D(scale=2)
>>> input = jnp.arange(1, 6).reshape(1, 5)
>>> print(layer(input))
[[1 1 2 2 3 3 4 4 5 5]]

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.Upsample2D(scale=1, method='nearest')

Upsample a 2D input to a given size using a given interpolation method.

Parameters:

• scale (int | tuple[int, ...]) – the scale of the output. accetps a single int or a tuple of two int denoting the scale multiplier along each axis.

• method (Literal['nearest', 'linear', 'cubic', 'lanczos3', 'lanczos5']) –

  Interpolation method defaults to nearest. choices are:

  • nearest: Nearest neighbor interpolation. The values of antialias and precision are ignored.

  • linear, bilinear, trilinear, triangle: Linear interpolation. If antialias is True, uses a triangular filter when downsampling.

  • cubic, bicubic, tricubic: Cubic interpolation, using the keys cubic kernel.

  • lanczos3: Lanczos resampling, using a kernel of radius 3.

  • lanczos5: Lanczos resampling, using a kernel of radius 5.

Example

>>> import serket as sk
>>> import jax.numpy as jnp
>>> layer = sk.nn.Upsample2D(scale=(1, 2))
>>> input = jnp.arange(1, 26).reshape(1, 5, 5)
>>> print(layer(input))
[[[ 1  1  2  2  3  3  4  4  5  5]
  [ 6  6  7  7  8  8  9  9 10 10]
  [11 11 12 12 13 13 14 14 15 15]
  [16 16 17 17 18 18 19 19 20 20]
  [21 21 22 22 23 23 24 24 25 25]]]

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

class serket.nn.Upsample3D(scale=1, method='nearest')

Upsample a 1D input to a given size using a given interpolation method.

Parameters:

• scale (int | tuple[int, ...]) – the scale of the output. accetps a single int or a tuple of three int denoting the scale multiplier along each axis.

• method (Literal['nearest', 'linear', 'cubic', 'lanczos3', 'lanczos5']) –

  Interpolation method defaults to nearest. choices are:

  • nearest: Nearest neighbor interpolation. The values of antialias and precision are ignored.

  • linear, bilinear, trilinear, triangle: Linear interpolation. If antialias is True, uses a triangular filter when downsampling.

  • cubic, bicubic, tricubic: Cubic interpolation, using the keys cubic kernel.

  • lanczos3: Lanczos resampling, using a kernel of radius 3.

  • lanczos5: Lanczos resampling, using a kernel of radius 5.

Example

>>> import serket as sk
>>> import jax.numpy as jnp
>>> layer = sk.nn.Upsample3D(scale=(1, 2, 1))
>>> input = jnp.arange(1, 9).reshape(1, 2, 2, 2)
>>> print(layer(input))
[[[[1 2]
  [1 2]
  [3 4]
  [3 4]]

 [[5 6]
  [5 6]
  [7 8]
  [7 8]]]]

__call__(input)

Call self as a function.

Parameters:

input (Array)

Return type:

Array

serket.nn.center_crop_nd(input, sizes)

Crops an input to the given size at the center.

Parameters:

• input (Array) – input array.

• sizes (tuple[int, ...]) – size of the crop along each axis.Accepts a tuple of int.

Return type:

Array

serket.nn.extract_patches(input, kernel_size, strides=1, padding='same')

Extract patches from an input array

Parameters:

• input (Array) – input array

• kernel_size (Union[int, Sequence[int]]) –

  Size of the convolutional kernel. accepts:

  • single integer for same kernel size in all dimensions.

  • sequence of integers for different kernel sizes in each dimension.

• strides (Union[int, Sequence[int]]) –

  stride of the convolution. Defaults to 1. accepts:

  • single integer for same stride in all dimensions.

  • sequence of integers for different strides in each dimension.

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

  Padding of the input before convolution. Default same. accepts:

  • single integer for same padding in all dimensions.

  • tuple of integers for different padding in each dimension.

  • tuple of a tuple of two integers for before and after padding in each dimension.

  • same/SAME for padding such that the output has the same shape as the input.

  • valid/VALID for no padding.

Returns:

A patches of the input array stacked along the first dimension.

Example

>>> import serket as sk
>>> import jax.numpy as jnp
>>> input = jnp.arange(15).reshape(5, 3)
>>> print(input)
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]
 [12 13 14]]
>>> kernel_size = 3
>>> strides = 1
>>> padding = "same"
>>> patches = sk.nn.extract_patches(input, kernel_size, strides, padding)
>>> print(patches.shape)
(15, 3, 3)
>>> print(patches[0])
[[0 0 0]
 [0 0 1]
 [0 3 4]]
>>> print(patches[1])
[[0 0 0]
 [0 1 2]
 [3 4 5]]

serket.nn.random_crop_nd(key, input, crop_size)

Crops an input to the given size at a random starts along each axis.

Parameters:

• key (Array) – random key.

• input (Array) – input array.

• crop_size (Sequence[int]) – size of the crop along each axis.Accepts a tuple of int.

Return type:

Array

serket.nn.upsample_nd(input, scale, method='nearest')

Upsample a 1D input to a given size using a given interpolation method.

Parameters:

• input (Array) – input array.

• scale (Union[int, Sequence[int]]) – the scale of the output. accetps a sequence of int denoting the scale multiplier along each axis.

• method (Literal['nearest', 'linear', 'cubic', 'lanczos3', 'lanczos5']) –

  Interpolation method defaults to nearest. choices are:

  • nearest: Nearest neighbor interpolation. The values of antialias and precision are ignored.

  • linear, bilinear, trilinear, triangle: Linear interpolation. If antialias is True, uses a triangular filter when downsampling.

  • cubic, bicubic, tricubic: Cubic interpolation, using the keys cubic kernel.

  • lanczos3: Lanczos resampling, using a kernel of radius 3.

  • lanczos5: Lanczos resampling, using a kernel of radius 5.

Return type:

Array