Example #1
0
class Sparse(Transform):
    """A sparse matrix transformation between an input and output signal.

    .. versionadded:: 3.0.0

    Parameters
    ----------
    shape : tuple of int
        The full shape of the sparse matrix: ``(size_out, size_in)``.
    indices : array_like of int
        An Nx2 array of integers indicating the (row,col) coordinates for the
        N non-zero elements in the matrix.
    init : `.Distribution` or array_like, optional
        A Distribution used to initialize the transform matrix, or a concrete
        instantiation for the matrix. If the matrix is square we also allow a
        scalar (equivalent to ``np.eye(n) * init``) or a vector (equivalent to
        ``np.diag(init)``) to represent the matrix more compactly.
    """

    shape = ShapeParam("shape", length=2, low=1)
    init = SparseInitParam("init")

    def __init__(self, shape, indices=None, init=1.0):
        super().__init__()

        self.shape = shape

        if scipy_sparse and isinstance(init, scipy_sparse.spmatrix):
            assert indices is None
            assert init.shape == self.shape
            self.init = init
        elif indices is not None:
            self.init = SparseMatrix(indices, init, shape)
        else:
            raise ValidationError(
                "Either `init` must be a `scipy.sparse.spmatrix`, "
                "or `indices` must be specified.",
                attr="init",
            )

    @property
    def _argreprs(self):
        return ["shape=%r" % (self.shape, )]

    def sample(self, rng=np.random):
        if scipy_sparse and isinstance(self.init, scipy_sparse.spmatrix):
            return self.init
        else:
            return self.init.sample(rng=rng)

    @property
    def size_in(self):
        return self.shape[1]

    @property
    def size_out(self):
        return self.shape[0]
Example #2
0
class Dense(Transform):
    """A dense matrix transformation between an input and output signal.

    .. versionadded:: 3.0.0

    Parameters
    ----------
    shape : tuple of int
        The shape of the dense matrix: ``(size_out, size_in)``.
    init : `.Distribution` or array_like, optional
        A Distribution used to initialize the transform matrix, or a concrete
        instantiation for the matrix. If the matrix is square we also allow a
        scalar (equivalent to ``np.eye(n) * init``) or a vector (equivalent to
        ``np.diag(init)``) to represent the matrix more compactly.
    """

    shape = ShapeParam("shape", length=2, low=1)
    init = DistOrArrayParam("init")

    def __init__(self, shape, init=1.0):
        super().__init__()

        self.shape = shape

        if is_array_like(init):
            init = np.asarray(init, dtype=rc.float_dtype)

            # check that the shape of init is compatible with the given shape
            # for this transform
            expected_shape = None
            if shape[0] != shape[1]:
                # init must be 2D if transform is not square
                expected_shape = shape
            elif init.ndim == 1:
                expected_shape = (shape[0], )
            elif init.ndim >= 2:
                expected_shape = shape

            if expected_shape is not None and init.shape != expected_shape:
                raise ValidationError(
                    "Shape of initial value %s does not match expected "
                    "shape %s" % (init.shape, expected_shape),
                    attr="init",
                )

        self.init = init

    @property
    def _argreprs(self):
        return ["shape=%r" % (self.shape, )]

    def sample(self, rng=np.random):
        if isinstance(self.init, Distribution):
            return self.init.sample(*self.shape, rng=rng)

        return self.init

    @property
    def init_shape(self):
        """The shape of the initial value."""
        return self.shape if isinstance(self.init,
                                        Distribution) else self.init.shape

    @property
    def size_in(self):
        return self.shape[1]

    @property
    def size_out(self):
        return self.shape[0]
Example #3
0
class Convolution(Transform):
    """An N-dimensional convolutional transform.

    The dimensionality of the convolution is determined by the input shape.

    .. versionadded:: 3.0.0

    Parameters
    ----------
    n_filters : int
        The number of convolutional filters to apply
    input_shape : tuple of int or `.ChannelShape`
        Shape of the input signal to the convolution; e.g.,
        ``(height, width, channels)`` for a 2D convolution with
        ``channels_last=True``.
    kernel_size : tuple of int, optional
        Size of the convolutional kernels (1 element for a 1D convolution,
        2 for a 2D convolution, etc.).
    strides : tuple of int, optional
        Stride of the convolution (1 element for a 1D convolution, 2 for
        a 2D convolution, etc.).
    padding : ``"same"`` or ``"valid"``, optional
        Padding method for input signal. "Valid" means no padding, and
        convolution will only be applied to the fully-overlapping areas of the
        input signal (meaning the output will be smaller). "Same" means that
        the input signal is zero-padded so that the output is the same shape
        as the input.
    channels_last : bool, optional
        If ``True`` (default), the channels are the last dimension in the input
        signal (e.g., a 28x28 image with 3 channels would have shape
        ``(28, 28, 3)``).  ``False`` means that channels are the first
        dimension (e.g., ``(3, 28, 28)``).
    init : `.Distribution` or `~numpy:numpy.ndarray`, optional
        A predefined kernel with shape
        ``kernel_size + (input_channels, n_filters)``, or a ``Distribution``
        that will be used to initialize the kernel.

    Notes
    -----
    As is typical in neural networks, this is technically correlation rather
    than convolution (because the kernel is not flipped).
    """

    n_filters = IntParam("n_filters", low=1)
    input_shape = ChannelShapeParam("input_shape", low=1)
    kernel_size = ShapeParam("kernel_size", low=1)
    strides = ShapeParam("strides", low=1)
    padding = EnumParam("padding", values=("same", "valid"))
    channels_last = BoolParam("channels_last")
    init = DistOrArrayParam("init")

    _param_init_order = ["channels_last", "input_shape"]

    def __init__(
            self,
            n_filters,
            input_shape,
            kernel_size=(3, 3),
            strides=(1, 1),
            padding="valid",
            channels_last=True,
            init=Uniform(-1, 1),
    ):
        super().__init__()

        self.n_filters = n_filters
        self.channels_last = channels_last  # must be set before input_shape
        self.input_shape = input_shape
        self.kernel_size = kernel_size
        self.strides = strides
        self.padding = padding
        self.init = init

        if len(kernel_size) != self.dimensions:
            raise ValidationError(
                "Kernel dimensions (%d) do not match input dimensions (%d)" %
                (len(kernel_size), self.dimensions),
                attr="kernel_size",
            )
        if len(strides) != self.dimensions:
            raise ValidationError(
                "Stride dimensions (%d) do not match input dimensions (%d)" %
                (len(strides), self.dimensions),
                attr="strides",
            )
        if not isinstance(init, Distribution):
            if init.shape != self.kernel_shape:
                raise ValidationError(
                    "Kernel shape %s does not match expected shape %s" %
                    (init.shape, self.kernel_shape),
                    attr="init",
                )

    @property
    def _argreprs(self):
        argreprs = [
            "n_filters=%r" % (self.n_filters, ),
            "input_shape=%s" % (self.input_shape.shape, ),
        ]
        if self.kernel_size != (3, 3):
            argreprs.append("kernel_size=%r" % (self.kernel_size, ))
        if self.strides != (1, 1):
            argreprs.append("strides=%r" % (self.strides, ))
        if self.padding != "valid":
            argreprs.append("padding=%r" % (self.padding, ))
        if self.channels_last is not True:
            argreprs.append("channels_last=%r" % (self.channels_last, ))
        return argreprs

    def sample(self, rng=np.random):
        if isinstance(self.init, Distribution):
            # we sample this way so that any variancescaling distribution based
            # on n/d is scaled appropriately
            kernel = [
                self.init.sample(self.input_shape.n_channels,
                                 self.n_filters,
                                 rng=rng)
                for _ in range(np.prod(self.kernel_size))
            ]
            kernel = np.reshape(kernel, self.kernel_shape)
        else:
            kernel = np.array(self.init, dtype=rc.float_dtype)
        return kernel

    @property
    def kernel_shape(self):
        """Full shape of kernel."""
        return self.kernel_size + (self.input_shape.n_channels, self.n_filters)

    @property
    def size_in(self):
        return self.input_shape.size

    @property
    def size_out(self):
        return self.output_shape.size

    @property
    def dimensions(self):
        """Dimensionality of convolution."""
        return self.input_shape.dimensions

    @property
    def output_shape(self):
        """Output shape after applying convolution to input."""
        output_shape = np.array(self.input_shape.spatial_shape,
                                dtype=rc.float_dtype)
        if self.padding == "valid":
            output_shape -= self.kernel_size
            output_shape += 1
        output_shape /= self.strides
        output_shape = tuple(np.ceil(output_shape).astype(rc.int_dtype))
        output_shape = (output_shape +
                        (self.n_filters, ) if self.channels_last else
                        (self.n_filters, ) + output_shape)

        return ChannelShape(output_shape, channels_last=self.channels_last)
Example #4
0
class SparseMatrix(FrozenObject):
    """Represents a sparse matrix.

    .. versionadded:: 3.0.0

    Parameters
    ----------
    indices : array_like of int
        An Nx2 array of integers indicating the (row,col) coordinates for the
        N non-zero elements in the matrix.
    data : array_like or `.Distribution`
        An Nx1 array defining the value of the nonzero elements in the matrix
        (corresponding to ``indices``), or a `.Distribution` that will be
        used to initialize the nonzero elements.
    shape : tuple of int
        Shape of the full matrix.
    """

    indices = NdarrayParam("indices", shape=("*", 2), dtype=np.int64)
    data = DistOrArrayParam("data", sample_shape=("*", ))
    shape = ShapeParam("shape", length=2)

    def __init__(self, indices, data, shape):
        super().__init__()

        self.indices = indices
        self.shape = shape

        # if data is not a distribution
        if is_array_like(data):
            data = np.asarray(data)

            # convert scalars to vectors
            if data.size == 1:
                data = data.item() * np.ones(self.indices.shape[0],
                                             dtype=data.dtype)

            if data.ndim != 1 or data.shape[0] != self.indices.shape[0]:
                raise ValidationError(
                    "Must be a vector of the same length as `indices`",
                    attr="data",
                    obj=self,
                )

        self.data = data
        self._allocated = None
        self._dense = None

    @property
    def dtype(self):
        return self.data.dtype

    @property
    def ndim(self):
        return len(self.shape)

    @property
    def size(self):
        return self.indices.shape[0]

    def allocate(self):
        """Return a `scipy.sparse.csr_matrix` or dense matrix equivalent.

        We mark this data as readonly to be consistent with how other
        data associated with signals are allocated. If this allocated
        data is to be modified, it should be copied first.
        """

        if self._allocated is not None:
            return self._allocated

        if scipy_sparse is None:
            warnings.warn("Sparse operations require Scipy, which is not "
                          "installed. Using dense matrices instead.")
            self._allocated = self.toarray().view()
        else:
            self._allocated = scipy_sparse.csr_matrix(
                (self.data, self.indices.T), shape=self.shape)
            self._allocated.data.setflags(write=False)

        return self._allocated

    def sample(self, rng=np.random):
        """Convert `.Distribution` data to fixed array.

        Parameters
        ----------
        rng : `.numpy.random.mtrand.RandomState`
            Random number generator that will be used when
            sampling distribution.

        Returns
        -------
        matrix : `.SparseMatrix`
            A new `.SparseMatrix` instance with `.Distribution` converted to
            array if ``self.data`` is a `.Distribution`, otherwise simply
            returns ``self``.
        """
        if isinstance(self.data, Distribution):
            return SparseMatrix(
                self.indices,
                self.data.sample(self.indices.shape[0], rng=rng),
                self.shape,
            )
        else:
            return self

    def toarray(self):
        """Return the dense matrix equivalent of this matrix."""

        if self._dense is not None:
            return self._dense

        self._dense = np.zeros(self.shape, dtype=self.dtype)
        self._dense[self.indices[:, 0], self.indices[:, 1]] = self.data
        # Mark as readonly, if the user wants to modify they should copy first
        self._dense.setflags(write=False)
        return self._dense
Example #5
0
class TensorNode(Node):
    """
    Inserts TensorFlow code into a Nengo model.

    Parameters
    ----------
    tensor_func : callable
        A function that maps node inputs to outputs
    shape_in : tuple of int
        Shape of TensorNode input signal (not including batch dimension).
    shape_out : tuple of int
        Shape of TensorNode output signal (not including batch dimension).
        If None, value will be inferred by calling ``tensor_func``.
    pass_time : bool
        If True, pass current simulation time to TensorNode function (in addition
        to the standard input).
    label : str (Default: None)
        A name for the node, used for debugging and visualization
    """

    tensor_func = TensorFuncParam("tensor_func")
    shape_in = ShapeParam("shape_in", default=None, low=1, optional=True)
    shape_out = ShapeParam("shape_out", default=None, low=1, optional=True)
    pass_time = BoolParam("pass_time", default=True)

    def __init__(
        self,
        tensor_func,
        shape_in=Default,
        shape_out=Default,
        pass_time=Default,
        label=Default,
    ):
        # pylint: disable=non-parent-init-called,super-init-not-called
        # note: we bypass the Node constructor, because we don't want to
        # perform validation on `output`
        NengoObject.__init__(self, label=label, seed=None)

        self.shape_in = shape_in
        self.shape_out = shape_out
        self.pass_time = pass_time

        if not (self.shape_in or self.pass_time):
            raise ValidationError(
                "Must specify either shape_in or pass_time", "TensorNode"
            )

        self.tensor_func = tensor_func

    @property
    def output(self):
        """
        Ensures that nothing tries to evaluate the `output` attribute
        (indicating that something is trying to simulate this as a regular
        `nengo.Node` rather than a TensorNode).
        """

        def output_func(*_):
            raise SimulationError(
                "Cannot call TensorNode output function (this probably means "
                "you are trying to use a TensorNode inside a Simulator other "
                "than NengoDL)"
            )

        return output_func

    @property
    def size_in(self):
        """Number of input elements (flattened)."""

        return 0 if self.shape_in is None else np.prod(self.shape_in)

    @property
    def size_out(self):
        """Number of output elements (flattened)."""

        return 0 if self.shape_out is None else np.prod(self.shape_out)
Example #6
0
class _ConvolutionBase(Transform):
    """Abstract base class for Convolution and ConvolutionTranspose."""

    n_filters = IntParam("n_filters", low=1)
    input_shape = ChannelShapeParam("input_shape", low=1)
    kernel_size = ShapeParam("kernel_size", low=1)
    strides = ShapeParam("strides", low=1)
    padding = EnumParam("padding", values=("same", "valid"))
    channels_last = BoolParam("channels_last")
    init = DistOrArrayParam("init")
    groups = IntParam("groups", low=1)

    _param_init_order = ["channels_last", "input_shape"]

    def __init__(
        self,
        n_filters,
        input_shape,
        kernel_size=(3, 3),
        strides=(1, 1),
        padding="valid",
        channels_last=True,
        init=Uniform(-1, 1),
        groups=1,
    ):
        super().__init__()

        self.n_filters = n_filters
        self.channels_last = channels_last  # must be set before input_shape
        self.input_shape = input_shape
        self.kernel_size = kernel_size
        self.strides = strides
        self.padding = padding
        self.init = init
        self.groups = groups

        if len(kernel_size) != self.dimensions:
            raise ValidationError(
                f"Kernel dimensions ({len(kernel_size)}) does not match "
                f"input dimensions ({self.dimensions})",
                attr="kernel_size",
            )
        if len(strides) != self.dimensions:
            raise ValidationError(
                f"Stride dimensions ({len(strides)}) does not match "
                f"input dimensions ({self.dimensions})",
                attr="strides",
            )
        if not isinstance(init, Distribution):
            if init.shape != self.kernel_shape:
                raise ValidationError(
                    f"Kernel shape {init.shape} does not match "
                    f"expected shape {self.kernel_shape}",
                    attr="init",
                )

        in_channels = self.input_shape.n_channels
        if groups > in_channels:
            raise ValidationError(
                f"Groups ({groups}) cannot be greater than "
                f"the number of input channels ({in_channels})",
                attr="groups",
            )
        if in_channels % groups != 0 or self.n_filters % groups != 0:
            raise ValidationError(
                f"Both the number of input channels ({in_channels}) and filters "
                f"({self.n_filters}) must be evenly divisible by ``groups`` ({groups})",
                attr="groups",
            )

    @property
    def _argreprs(self):
        argreprs = [
            f"n_filters={self.n_filters!r}",
            f"input_shape={self.input_shape.shape}",
        ]
        if self.kernel_size != (3, 3):
            argreprs.append(f"kernel_size={self.kernel_size!r}")
        if self.strides != (1, 1):
            argreprs.append(f"strides={self.strides!r}")
        if self.padding != "valid":
            argreprs.append(f"padding={self.padding!r}")
        if self.channels_last is not True:
            argreprs.append(f"channels_last={self.channels_last!r}")
        if self.groups != 1:
            argreprs.append(f"groups={self.groups!r}")
        return argreprs

    def sample(self, rng=np.random):
        if isinstance(self.init, Distribution):
            # we sample this way so that any variancescaling distribution based
            # on n/d is scaled appropriately
            kernel = [
                self.init.sample(
                    self.input_shape.n_channels // self.groups, self.n_filters, rng=rng
                )
                for _ in range(np.prod(self.kernel_size))
            ]
            kernel = np.reshape(kernel, self.kernel_shape)
        else:
            kernel = np.array(self.init, dtype=rc.float_dtype)
        return kernel

    @property
    def kernel_shape(self):
        """Full shape of kernel."""
        return self.kernel_size + (
            self.input_shape.n_channels // self.groups,
            self.n_filters,
        )

    @property
    def size_in(self):
        return self.input_shape.size

    @property
    def size_out(self):
        return self.output_shape.size

    @property
    def dimensions(self):
        """Dimensionality of convolution."""
        return self.input_shape.dimensions

    def _forward_shape(self, input_spatial_shape, n_filters):
        output_shape = np.array(input_spatial_shape, dtype=rc.float_dtype)
        if self.padding == "valid":
            output_shape -= self.kernel_size
            output_shape += 1
        output_shape /= self.strides
        output_shape = tuple(np.ceil(output_shape).astype(rc.int_dtype))

        return ChannelShape.from_space_and_channels(
            output_shape, n_filters, channels_last=self.channels_last
        )
Example #7
0
class PresentJitteredImages(Process):
    images = NdarrayParam('images', shape=('...', ))
    image_shape = ShapeParam('image_shape', length=3, low=1)
    output_shape = ShapeParam('output_shape', length=2, low=1)
    presentation_time = NumberParam('presentation_time', low=0, low_open=True)
    jitter_std = NumberParam('jitter_std', low=0, low_open=True, optional=True)
    jitter_tau = NumberParam('jitter_tau', low=0, low_open=True)

    def __init__(self,
                 images,
                 presentation_time,
                 output_shape,
                 jitter_std=None,
                 jitter_tau=None,
                 **kwargs):
        import scipy.ndimage.interpolation
        # ^ required for simulation, so check it here

        self.images = images
        self.presentation_time = presentation_time
        self.image_shape = images.shape[1:]
        self.output_shape = output_shape
        self.jitter_std = jitter_std
        self.jitter_tau = (presentation_time
                           if jitter_tau is None else jitter_tau)

        nc = self.image_shape[0]
        nyi, nyj = self.output_shape
        super(PresentJitteredImages,
              self).__init__(default_size_in=0,
                             default_size_out=nc * nyi * nyj,
                             **kwargs)

    def make_step(self, shape_in, shape_out, dt, rng):
        import scipy.ndimage.interpolation

        nc, nxi, nxj = self.image_shape
        nyi, nyj = self.output_shape
        ni, nj = nxi - nyi, nxj - nyj
        nij = np.array([ni, nj])
        assert shape_in == (0, )
        assert shape_out == (nc * nyi * nyj, )

        if self.jitter_std is None:
            si, sj = ni / 4., nj / 4.
        else:
            si = sj = self.jitter_std

        tau = self.jitter_tau

        n = len(self.images)
        images = self.images.reshape((n, nc, nxi, nxj))
        presentation_time = float(self.presentation_time)

        cij = (nij - 1) / 2.
        dt7tau = dt / tau
        sigma2 = np.sqrt(2. * dt / tau) * np.array([si, sj])
        ij = cij.copy()

        def step_presentjitteredimages(t):
            # update jitter position
            ij0 = dt7tau * (cij - ij) + sigma2 * rng.normal(size=2)
            ij[:] = (ij + ij0).clip((0, 0), (ni, nj))

            # select image
            k = int((t - dt) / presentation_time + 1e-7)
            image = images[k % n]

            # interpolate jittered sub-image
            i, j = ij
            image = scipy.ndimage.interpolation.shift(
                image, (0, ni - i, nj - j))[:, -nyi:, -nyj:]

            return image.ravel()

        return step_presentjitteredimages
Example #8
0
class Pool2d(Process):
    """Perform 2-D (image) pooling on an input.

    Parameters
    ----------
    shape_in : 3-tuple (channels, height, width)
        Shape of the input image.
    pool_size : 2-tuple (vertical, horizontal) or int
        Shape of the pooling region. If an integer is provided, the shape will
        be square with the given side length.
    strides : 2-tuple (vertical, horizontal) or int
        Spacing between pooling placements. If ``None`` (default), will be
        equal to ``pool_size`` resulting in non-overlapping pooling.
    kind : "avg" or "max"
        Type of pooling to perform: average pooling or max pooling.
    mode : "full" or "valid"
        If the input image does not divide into an integer number of pooling
        regions, whether to add partial pooling regions for the extra
        pixels ("full"), or discard extra input pixels ("valid").

    Attributes
    ----------
    shape_out : 3-tuple (channels, height, width)
        Shape of the output image.
    """
    shape_in = ShapeParam('shape_in', length=3, low=1)
    shape_out = ShapeParam('shape_out', length=3, low=1)
    pool_size = ShapeParam('pool_size', length=2, low=1)
    strides = ShapeParam('strides', length=2, low=1)
    kind = EnumParam('kind', values=('avg', 'max'))
    mode = EnumParam('mode', values=('full', 'valid'))

    def __init__(self,
                 shape_in,
                 pool_size,
                 strides=None,
                 kind='avg',
                 mode='full'):
        self.shape_in = shape_in
        self.pool_size = (pool_size
                          if is_iterable(pool_size) else [pool_size] * 2)
        self.strides = (strides if is_iterable(strides) else [strides] *
                        2 if strides is not None else self.pool_size)
        self.kind = kind
        self.mode = mode
        if not all(st <= p for st, p in zip(self.strides, self.pool_size)):
            raise ValueError("Strides %s must be <= pool_size %s" %
                             (self.strides, self.pool_size))

        nc, nxi, nxj = self.shape_in
        nyi_float = float(nxi - self.pool_size[0]) / self.strides[0]
        nyj_float = float(nxj - self.pool_size[1]) / self.strides[1]
        if self.mode == 'full':
            nyi = 1 + int(np.ceil(nyi_float))
            nyj = 1 + int(np.ceil(nyj_float))
        elif self.mode == 'valid':
            nyi = 1 + int(np.floor(nyi_float))
            nyj = 1 + int(np.floor(nyj_float))
        self.shape_out = (nc, nyi, nyj)

        super(Pool2d, self).__init__(default_size_in=np.prod(self.shape_in),
                                     default_size_out=np.prod(self.shape_out))

    def make_step(self, shape_in, shape_out, dt, rng):
        assert np.prod(shape_in) == np.prod(self.shape_in)
        assert np.prod(shape_out) == np.prod(self.shape_out)
        nc, nxi, nxj = self.shape_in
        nc, nyi, nyj = self.shape_out
        si, sj = self.pool_size
        sti, stj = self.strides
        kind = self.kind
        nxi2, nxj2 = nyi * sti, nyj * stj

        def step_pool2d(t, x):
            x = x.reshape(-1, nc, nxi, nxj)
            y = np.zeros((x.shape[0], nc, nyi, nyj), dtype=x.dtype)
            n = np.zeros((nyi, nyj))

            for i in range(si):
                for j in range(sj):
                    xij = x[:, :, i:min(nxi2 + i, nxi):sti,
                            j:min(nxj2 + j, nxj):stj]
                    ni, nj = xij.shape[-2:]
                    if kind == 'max':
                        y[:, :, :ni, :nj] = np.maximum(y[:, :, :ni, :nj], xij)
                    elif kind == 'avg':
                        y[:, :, :ni, :nj] += xij
                        n[:ni, :nj] += 1
                    else:
                        raise NotImplementedError(kind)

            if kind == 'avg':
                y /= n

            return y.ravel()

        return step_pool2d
Example #9
0
class Conv2d(Process):
    """Perform 2-D (image) convolution on an input.

    Parameters
    ----------
    shape_in : 3-tuple (n_channels, height, width)
        Shape of the input images: channels, height, width.
    filters : array_like (n_filters, n_channels, f_height, f_width)
        Static filters to convolve with the input. Shape is number of filters,
        number of input channels, filter height, and filter width. Shape can
        also be (n_filters, height, width, n_channels, f_height, f_width)
        to apply different filters at each point in the image, where 'height'
        and 'width' are the input image height and width.
    biases : array_like (1,) or (n_filters,) or (n_filters, height, width)
        Biases to add to outputs. Can have one bias across the entire output
        space, one bias per filter, or a unique bias for each output pixel.
    strides : 2-tuple (vertical, horizontal) or int
        Spacing between filter placements. If an integer
        is provided, the same spacing is used in both dimensions.
    padding : 2-tuple (vertical, horizontal) or int
        Amount of zero-padding around the outside of the input image. Padding
        is applied to both sides, e.g. ``padding=(1, 0)`` will add one pixel
        of padding to the top and bottom, and none to the left and right.
    """

    shape_in = ShapeParam('shape_in', length=3, low=1)
    shape_out = ShapeParam('shape_out', length=3, low=1)
    strides = ShapeParam('strides', length=2, low=1)
    padding = ShapeParam('padding', length=2)
    filters = NdarrayParam('filters', shape=('...', ))
    biases = NdarrayParam('biases', shape=('...', ), optional=True)
    border = EnumParam('border', values=('floor', 'ceil'))

    def __init__(self,
                 shape_in,
                 filters,
                 biases=None,
                 strides=1,
                 padding=0,
                 border='ceil'):  # noqa: C901
        self.shape_in = shape_in
        self.filters = filters
        if self.filters.ndim not in [4, 6]:
            raise ValueError(
                "`filters` must have four or six dimensions "
                "(filters, [height, width,] channels, f_height, f_width)")
        if self.filters.shape[-3] != self.shape_in[0]:
            raise ValueError(
                "Filter channels (%d) and input channels (%d) must match" %
                (self.filters.shape[-3], self.shape_in[0]))
        if not all(s % 2 == 1 for s in self.filters.shape[-2:]):
            raise ValueError("Filter shapes must be odd (got %r)" %
                             (self.filters.shape[-2:], ))

        self.strides = strides if is_iterable(strides) else [strides] * 2
        self.padding = padding if is_iterable(padding) else [padding] * 2
        self.border = border

        nf = self.filters.shape[0]
        nxi, nxj = self.shape_in[1:]
        si, sj = self.filters.shape[-2:]
        pi, pj = self.padding
        sti, stj = self.strides
        rounder = np.ceil if self.border == 'ceil' else np.floor
        nyi = 1 + max(int(rounder(float(2 * pi + nxi - si) / sti)), 0)
        nyj = 1 + max(int(rounder(float(2 * pj + nxj - sj) / stj)), 0)
        self.shape_out = (nf, nyi, nyj)
        if self.filters.ndim == 6 and self.filters.shape[1:3] != (nyi, nyj):
            raise ValueError("Number of local filters %r must match out shape "
                             "%r" % (self.filters.shape[1:3], (nyi, nyj)))

        self.biases = biases if biases is not None else None
        if self.biases is not None:
            if self.biases.size == 1:
                self.biases.shape = (1, 1, 1)
            elif self.biases.size == np.prod(self.shape_out):
                self.biases.shape = self.shape_out
            elif self.biases.size == self.shape_out[0]:
                self.biases.shape = (self.shape_out[0], 1, 1)
            elif self.biases.size == np.prod(self.shape_out[1:]):
                self.biases.shape = (1, ) + self.shape_out[1:]
            else:
                raise ValueError(
                    "Biases size (%d) does not match output shape %s" %
                    (self.biases.size, self.shape_out))

        super(Conv2d, self).__init__(default_size_in=np.prod(self.shape_in),
                                     default_size_out=np.prod(self.shape_out))

    def make_step(self, shape_in, shape_out, dt, rng):
        assert np.prod(shape_in) == np.prod(self.shape_in)
        assert np.prod(shape_out) == np.prod(self.shape_out)
        shape_in, shape_out = self.shape_in, self.shape_out

        filters = self.filters
        local_filters = filters.ndim == 6
        biases = self.biases

        nc, nxi, nxj = shape_in
        nf, nyi, nyj = shape_out
        si, sj = filters.shape[-2:]
        pi, pj = self.padding
        sti, stj = self.strides

        def step_conv2d(t, x):
            x = x.reshape(-1, nc, nxi, nxj)
            n = x.shape[0]
            y = np.zeros((n, nf, nyi, nyj), dtype=x.dtype)

            for i in range(nyi):
                for j in range(nyj):
                    i0 = i * sti - pi
                    j0 = j * stj - pj
                    i1, j1 = i0 + si, j0 + sj
                    sli = slice(max(-i0, 0), min(nxi + si - i1, si))
                    slj = slice(max(-j0, 0), min(nxj + sj - j1, sj))
                    w = (filters[:, i, j, :, sli,
                                 slj] if local_filters else filters[:, :, sli,
                                                                    slj])
                    xij = x[:, :,
                            max(i0, 0):min(i1, nxi),
                            max(j0, 0):min(j1, nxj)]
                    y[:, :, i, j] = np.dot(xij.reshape(n, -1),
                                           w.reshape(nf, -1).T)

            if biases is not None:
                y += biases

            return y.ravel()

        return step_conv2d