ReplicationPad3d

Contents

ReplicationPad3d#

class brainstate.nn.ReplicationPad3d(padding, in_size=None, name=None)[source]#

Pads the input tensor using replication of the input boundary.

Parameters:

  • padding (int | Sequence[int]) –

    The size of the padding. Can be:

    • int: same padding for all sides

    • Sequence[int] of length 3: (depth_pad, height_pad, width_pad)

    • Sequence[int] of length 6: (depth_before, depth_after, height_before, height_after, width_before, width_after)

    Note

    For the length-6 form, the pairs pad the spatial axes in order (depth, then height, then width). This order is the REVERSE of torch.nn.*Pad3d, which lists the last spatial axis (width) first as (left, right, top, bottom, front, back).

  • in_size (int | Sequence[int] | integer | Sequence[integer] | None) – The input size.

  • name (str | None) – The name of the module.

Examples

>>> import brainstate as brainstate
>>> import jax.numpy as jnp
>>> pad = brainstate.nn.ReplicationPad3d(1)
>>> input = jnp.ones((1, 4, 4, 4, 3))
>>> output = pad(input)
>>> print(output.shape)
(1, 6, 6, 6, 3)
Contents