Skip to content

Block Factories¤

pdequinox.blocks.ClassicDoubleConvBlockFactory ¤

Bases: BlockFactory

Source code in pdequinox/blocks/_classic_double_conv_block.py 
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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
class ClassicDoubleConvBlockFactory(BlockFactory):
    kernel_size: int
    use_norm: bool
    num_groups: int
    use_bias: bool
    zero_bias_init: bool

    def __init__(
        self,
        *,
        kernel_size: int = 3,
        use_norm: bool = True,
        num_groups: int = 1,
        use_bias: bool = True,
        zero_bias_init: bool = False,
    ):
        """
        Factory for creating `ClassicDoubleConvBlock` instances.

        **Arguments:**

        - `kernel_size`: The size of the convolutional kernel. Default is `3`.
        - `use_norm`: Whether to use group normalization. Default is `True`.
        - `num_groups`: The number of groups to use for group normalization.
            Default is `1`.
        - `use_bias`: Whether to use bias in the convolutional layers. Default is
            `True`.
        - `zero_bias_init`: Whether to initialise the bias to zero. Default is
            `False`.
        """
        self.kernel_size = kernel_size
        self.use_norm = use_norm
        self.num_groups = num_groups
        self.use_bias = use_bias
        self.zero_bias_init = zero_bias_init

    def __call__(
        self,
        num_spatial_dims: int,
        in_channels: int,
        out_channels: int,
        *,
        activation: Callable,
        boundary_mode: Literal["periodic", "dirichlet", "neumann"],
        key: PRNGKeyArray,
    ):
        return ClassicDoubleConvBlock(
            num_spatial_dims=num_spatial_dims,
            in_channels=in_channels,
            out_channels=out_channels,
            activation=activation,
            boundary_mode=boundary_mode,
            key=key,
            kernel_size=self.kernel_size,
            use_norm=self.use_norm,
            num_groups=self.num_groups,
            use_bias=self.use_bias,
            zero_bias_init=self.zero_bias_init,
        )
__init__ ¤
__init__(
    *,
    kernel_size: int = 3,
    use_norm: bool = True,
    num_groups: int = 1,
    use_bias: bool = True,
    zero_bias_init: bool = False
)

Factory for creating ClassicDoubleConvBlock instances.

Arguments:

  • kernel_size: The size of the convolutional kernel. Default is 3.
  • use_norm: Whether to use group normalization. Default is True.
  • num_groups: The number of groups to use for group normalization. Default is 1.
  • use_bias: Whether to use bias in the convolutional layers. Default is True.
  • zero_bias_init: Whether to initialise the bias to zero. Default is False.
Source code in pdequinox/blocks/_classic_double_conv_block.py 
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
def __init__(
    self,
    *,
    kernel_size: int = 3,
    use_norm: bool = True,
    num_groups: int = 1,
    use_bias: bool = True,
    zero_bias_init: bool = False,
):
    """
    Factory for creating `ClassicDoubleConvBlock` instances.

    **Arguments:**

    - `kernel_size`: The size of the convolutional kernel. Default is `3`.
    - `use_norm`: Whether to use group normalization. Default is `True`.
    - `num_groups`: The number of groups to use for group normalization.
        Default is `1`.
    - `use_bias`: Whether to use bias in the convolutional layers. Default is
        `True`.
    - `zero_bias_init`: Whether to initialise the bias to zero. Default is
        `False`.
    """
    self.kernel_size = kernel_size
    self.use_norm = use_norm
    self.num_groups = num_groups
    self.use_bias = use_bias
    self.zero_bias_init = zero_bias_init
__call__ ¤
__call__(
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal[
        "periodic", "dirichlet", "neumann"
    ],
    key: PRNGKeyArray
)
Source code in pdequinox/blocks/_classic_double_conv_block.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def __call__(
    self,
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal["periodic", "dirichlet", "neumann"],
    key: PRNGKeyArray,
):
    return ClassicDoubleConvBlock(
        num_spatial_dims=num_spatial_dims,
        in_channels=in_channels,
        out_channels=out_channels,
        activation=activation,
        boundary_mode=boundary_mode,
        key=key,
        kernel_size=self.kernel_size,
        use_norm=self.use_norm,
        num_groups=self.num_groups,
        use_bias=self.use_bias,
        zero_bias_init=self.zero_bias_init,
    )

pdequinox.blocks.ClassicResBlockFactory ¤

Bases: Module

Source code in pdequinox/blocks/_classic_res_block.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
class ClassicResBlockFactory(eqx.Module):
    kernel_size: int
    use_norm: bool
    num_groups: int
    use_bias: bool
    zero_bias_init: bool

    def __init__(
        self,
        kernel_size: int = 3,
        *,
        use_norm: bool = False,
        num_groups: int = 1,
        use_bias: bool = True,
        zero_bias_init: bool = False,
    ):
        """
        Factory for creating `ClassicResBlock` instances.

        **Arguments:**

        - `kernel_size`: The size of the convolutional kernel. Default is `3`.
        - `use_norm`: Whether to use group normalization. Default is `True`.
        - `num_groups`: The number of groups to use for group normalization.
            Default is `1`.
        - `use_bias`: Whether to use bias in the convolutional layers. Default is
            `True`.
        - `zero_bias_init`: Whether to initialise the bias to zero. Default is
            `False`.
        """
        self.kernel_size = kernel_size
        self.use_norm = use_norm
        self.num_groups = num_groups
        self.use_bias = use_bias
        self.zero_bias_init = zero_bias_init

    def __call__(
        self,
        num_spatial_dims: int,
        in_channels: int,
        out_channels: int,
        *,
        activation: Callable,
        boundary_mode: Literal["periodic", "dirichlet", "neumann"],
        key: PRNGKeyArray,
    ):
        return ClassicResBlock(
            num_spatial_dims,
            in_channels,
            out_channels,
            activation=activation,
            kernel_size=self.kernel_size,
            use_norm=self.use_norm,
            num_groups=self.num_groups,
            boundary_mode=boundary_mode,
            key=key,
            use_bias=self.use_bias,
            zero_bias_init=self.zero_bias_init,
        )
__init__ ¤
__init__(
    kernel_size: int = 3,
    *,
    use_norm: bool = False,
    num_groups: int = 1,
    use_bias: bool = True,
    zero_bias_init: bool = False
)

Factory for creating ClassicResBlock instances.

Arguments:

  • kernel_size: The size of the convolutional kernel. Default is 3.
  • use_norm: Whether to use group normalization. Default is True.
  • num_groups: The number of groups to use for group normalization. Default is 1.
  • use_bias: Whether to use bias in the convolutional layers. Default is True.
  • zero_bias_init: Whether to initialise the bias to zero. Default is False.
Source code in pdequinox/blocks/_classic_res_block.py 
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def __init__(
    self,
    kernel_size: int = 3,
    *,
    use_norm: bool = False,
    num_groups: int = 1,
    use_bias: bool = True,
    zero_bias_init: bool = False,
):
    """
    Factory for creating `ClassicResBlock` instances.

    **Arguments:**

    - `kernel_size`: The size of the convolutional kernel. Default is `3`.
    - `use_norm`: Whether to use group normalization. Default is `True`.
    - `num_groups`: The number of groups to use for group normalization.
        Default is `1`.
    - `use_bias`: Whether to use bias in the convolutional layers. Default is
        `True`.
    - `zero_bias_init`: Whether to initialise the bias to zero. Default is
        `False`.
    """
    self.kernel_size = kernel_size
    self.use_norm = use_norm
    self.num_groups = num_groups
    self.use_bias = use_bias
    self.zero_bias_init = zero_bias_init
__call__ ¤
__call__(
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal[
        "periodic", "dirichlet", "neumann"
    ],
    key: PRNGKeyArray
)
Source code in pdequinox/blocks/_classic_res_block.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def __call__(
    self,
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal["periodic", "dirichlet", "neumann"],
    key: PRNGKeyArray,
):
    return ClassicResBlock(
        num_spatial_dims,
        in_channels,
        out_channels,
        activation=activation,
        kernel_size=self.kernel_size,
        use_norm=self.use_norm,
        num_groups=self.num_groups,
        boundary_mode=boundary_mode,
        key=key,
        use_bias=self.use_bias,
        zero_bias_init=self.zero_bias_init,
    )

pdequinox.blocks.ModernResBlockFactory ¤

Bases: Module

Source code in pdequinox/blocks/_modern_res_block.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
class ModernResBlockFactory(eqx.Module):
    kernel_size: int
    use_norm: bool
    num_groups: int
    use_bias: bool
    zero_bias_init: bool

    def __init__(
        self,
        kernel_size: int = 3,
        *,
        use_norm: bool = True,
        num_groups: int = 1,
        use_bias: bool = True,
        zero_bias_init: bool = False,
    ):
        """
        Factory for creating `ModernResBlock` instances.

        **Arguments:**

        - `kernel_size`: The size of the convolutional kernel. Default is `3`.
        - `use_norm`: Whether to use group normalization. Default is `True`.
        - `num_groups`: The number of groups to use for group normalization.
            Default is `1`.
        - `use_bias`: Whether to use bias in the convolutional layers. Default is
            `True`.
        - `zero_bias_init`: Whether to initialise the bias to zero. Default is
            `False`.
        """
        self.kernel_size = kernel_size
        self.use_norm = use_norm
        self.num_groups = num_groups
        self.use_bias = use_bias
        self.zero_bias_init = zero_bias_init

    def __call__(
        self,
        num_spatial_dims: int,
        in_channels: int,
        out_channels: int,
        *,
        activation: Callable,
        boundary_mode: Literal["periodic", "dirichlet", "neumann"],
        key: PRNGKeyArray,
    ):
        return ModernResBlock(
            num_spatial_dims=num_spatial_dims,
            in_channels=in_channels,
            out_channels=out_channels,
            activation=activation,
            kernel_size=self.kernel_size,
            boundary_mode=boundary_mode,
            key=key,
            use_norm=self.use_norm,
            num_groups=self.num_groups,
            use_bias=self.use_bias,
            zero_bias_init=self.zero_bias_init,
        )
__init__ ¤
__init__(
    kernel_size: int = 3,
    *,
    use_norm: bool = True,
    num_groups: int = 1,
    use_bias: bool = True,
    zero_bias_init: bool = False
)

Factory for creating ModernResBlock instances.

Arguments:

  • kernel_size: The size of the convolutional kernel. Default is 3.
  • use_norm: Whether to use group normalization. Default is True.
  • num_groups: The number of groups to use for group normalization. Default is 1.
  • use_bias: Whether to use bias in the convolutional layers. Default is True.
  • zero_bias_init: Whether to initialise the bias to zero. Default is False.
Source code in pdequinox/blocks/_modern_res_block.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
def __init__(
    self,
    kernel_size: int = 3,
    *,
    use_norm: bool = True,
    num_groups: int = 1,
    use_bias: bool = True,
    zero_bias_init: bool = False,
):
    """
    Factory for creating `ModernResBlock` instances.

    **Arguments:**

    - `kernel_size`: The size of the convolutional kernel. Default is `3`.
    - `use_norm`: Whether to use group normalization. Default is `True`.
    - `num_groups`: The number of groups to use for group normalization.
        Default is `1`.
    - `use_bias`: Whether to use bias in the convolutional layers. Default is
        `True`.
    - `zero_bias_init`: Whether to initialise the bias to zero. Default is
        `False`.
    """
    self.kernel_size = kernel_size
    self.use_norm = use_norm
    self.num_groups = num_groups
    self.use_bias = use_bias
    self.zero_bias_init = zero_bias_init
__call__ ¤
__call__(
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal[
        "periodic", "dirichlet", "neumann"
    ],
    key: PRNGKeyArray
)
Source code in pdequinox/blocks/_modern_res_block.py 
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def __call__(
    self,
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal["periodic", "dirichlet", "neumann"],
    key: PRNGKeyArray,
):
    return ModernResBlock(
        num_spatial_dims=num_spatial_dims,
        in_channels=in_channels,
        out_channels=out_channels,
        activation=activation,
        kernel_size=self.kernel_size,
        boundary_mode=boundary_mode,
        key=key,
        use_norm=self.use_norm,
        num_groups=self.num_groups,
        use_bias=self.use_bias,
        zero_bias_init=self.zero_bias_init,
    )

pdequinox.blocks.ClassicSpectralBlockFactory ¤

Bases: BlockFactory

Source code in pdequinox/blocks/_classic_spectral_block.py 
 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
class ClassicSpectralBlockFactory(BlockFactory):
    num_modes: Union[int, tuple[int, ...]]
    use_bias: bool
    zero_bias_init: bool

    def __init__(
        self,
        *,
        num_modes: int = 8,
        use_bias: bool = True,
        zero_bias_init: bool = False,
    ):
        """
        Factory for creating `ClassicSpectralBlock` instances.

        **Arguments:**

        - `num_modes`: How many modes to consider in Fourier space. At max this
            can be N//2+1, with N being the number of spatial points. Think of
            it as the analogy of the kernel size.
        - `use_bias`: Whether to use a bias in the bypass convolution. Default
            `True`.
        - `zero_bias_init`: Whether to initialise the bias to zero. Default is
            `False`.
        """
        self.num_modes = num_modes
        self.use_bias = use_bias
        self.zero_bias_init = zero_bias_init

    def __call__(
        self,
        num_spatial_dims: int,
        in_channels: int,
        out_channels: int,
        *,
        activation: Callable,
        boundary_mode: Literal["periodic", "dirichlet", "neumann"],  # unused
        key: PRNGKeyArray,
        # unused
    ):
        return ClassicSpectralBlock(
            num_spatial_dims=num_spatial_dims,
            in_channels=in_channels,
            out_channels=out_channels,
            num_modes=self.num_modes,
            activation=activation,
            key=key,
            use_bias=self.use_bias,
            zero_bias_init=self.zero_bias_init,
        )
__init__ ¤
__init__(
    *,
    num_modes: int = 8,
    use_bias: bool = True,
    zero_bias_init: bool = False
)

Factory for creating ClassicSpectralBlock instances.

Arguments:

  • num_modes: How many modes to consider in Fourier space. At max this can be N//2+1, with N being the number of spatial points. Think of it as the analogy of the kernel size.
  • use_bias: Whether to use a bias in the bypass convolution. Default True.
  • zero_bias_init: Whether to initialise the bias to zero. Default is False.
Source code in pdequinox/blocks/_classic_spectral_block.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def __init__(
    self,
    *,
    num_modes: int = 8,
    use_bias: bool = True,
    zero_bias_init: bool = False,
):
    """
    Factory for creating `ClassicSpectralBlock` instances.

    **Arguments:**

    - `num_modes`: How many modes to consider in Fourier space. At max this
        can be N//2+1, with N being the number of spatial points. Think of
        it as the analogy of the kernel size.
    - `use_bias`: Whether to use a bias in the bypass convolution. Default
        `True`.
    - `zero_bias_init`: Whether to initialise the bias to zero. Default is
        `False`.
    """
    self.num_modes = num_modes
    self.use_bias = use_bias
    self.zero_bias_init = zero_bias_init
__call__ ¤
__call__(
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal[
        "periodic", "dirichlet", "neumann"
    ],
    key: PRNGKeyArray
)
Source code in pdequinox/blocks/_classic_spectral_block.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def __call__(
    self,
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal["periodic", "dirichlet", "neumann"],  # unused
    key: PRNGKeyArray,
    # unused
):
    return ClassicSpectralBlock(
        num_spatial_dims=num_spatial_dims,
        in_channels=in_channels,
        out_channels=out_channels,
        num_modes=self.num_modes,
        activation=activation,
        key=key,
        use_bias=self.use_bias,
        zero_bias_init=self.zero_bias_init,
    )

pdequinox.blocks.DilatedResBlockFactory ¤

Bases: Module

Source code in pdequinox/blocks/_dilated_res_block.py 
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
class DilatedResBlockFactory(eqx.Module):
    kernel_size: int
    dilation_rates: tuple[int]
    use_norm: bool
    num_groups: int
    use_bias: bool
    zero_bias_init: bool

    def __init__(
        self,
        kernel_size: int = 3,
        dilation_rates: tuple[int] = (1, 2, 4, 8, 4, 2, 1),
        *,
        use_norm: bool = True,
        num_groups: int = 1,
        use_bias: bool = True,
        zero_bias_init: bool = False,
    ):
        """
        Factory for creating `DilatedResBlock` instances.

        **Arguments:**

        - `kernel_size`: The size of the convolutional kernel. Default is `3`.
        - `dilation_rates`: A sequence of integers. Their length identifies the
            number of sequential convolutions performed. Each integer is the
            dilation performed at that convolution. Typically, this list follows
            the pattern of first increasing in dilation rate, and then
            decreasing again. Default is `(1, 2, 4, 8, 4, 2, 1)`.
        - `use_norm`: Whether to use group normalization. Default is `True`.
        - `num_groups`: The number of groups to use for group normalization.
            Default is `1`.
        - `use_bias`: Whether to use bias in the convolutional layers. Default is
            `True`.
        - `zero_bias_init`: Whether to initialise the bias to zero. Default is
            `False`.
        """
        self.kernel_size = kernel_size
        self.dilation_rates = dilation_rates
        self.use_norm = use_norm
        self.num_groups = num_groups
        self.use_bias = use_bias
        self.zero_bias_init = zero_bias_init

    def __call__(
        self,
        num_spatial_dims: int,
        in_channels: int,
        out_channels: int,
        *,
        activation: Callable,
        boundary_mode: Literal["periodic", "dirichlet", "neumann"],
        key: PRNGKeyArray,
    ) -> DilatedResBlock:
        return DilatedResBlock(
            num_spatial_dims=num_spatial_dims,
            in_channels=in_channels,
            out_channels=out_channels,
            activation=activation,
            kernel_size=self.kernel_size,
            dilation_rates=self.dilation_rates,
            boundary_mode=boundary_mode,
            key=key,
            use_norm=self.use_norm,
            num_groups=self.num_groups,
            use_bias=self.use_bias,
            zero_bias_init=self.zero_bias_init,
        )
__init__ ¤
__init__(
    kernel_size: int = 3,
    dilation_rates: tuple[int] = (1, 2, 4, 8, 4, 2, 1),
    *,
    use_norm: bool = True,
    num_groups: int = 1,
    use_bias: bool = True,
    zero_bias_init: bool = False
)

Factory for creating DilatedResBlock instances.

Arguments:

  • kernel_size: The size of the convolutional kernel. Default is 3.
  • dilation_rates: A sequence of integers. Their length identifies the number of sequential convolutions performed. Each integer is the dilation performed at that convolution. Typically, this list follows the pattern of first increasing in dilation rate, and then decreasing again. Default is (1, 2, 4, 8, 4, 2, 1).
  • use_norm: Whether to use group normalization. Default is True.
  • num_groups: The number of groups to use for group normalization. Default is 1.
  • use_bias: Whether to use bias in the convolutional layers. Default is True.
  • zero_bias_init: Whether to initialise the bias to zero. Default is False.
Source code in pdequinox/blocks/_dilated_res_block.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
def __init__(
    self,
    kernel_size: int = 3,
    dilation_rates: tuple[int] = (1, 2, 4, 8, 4, 2, 1),
    *,
    use_norm: bool = True,
    num_groups: int = 1,
    use_bias: bool = True,
    zero_bias_init: bool = False,
):
    """
    Factory for creating `DilatedResBlock` instances.

    **Arguments:**

    - `kernel_size`: The size of the convolutional kernel. Default is `3`.
    - `dilation_rates`: A sequence of integers. Their length identifies the
        number of sequential convolutions performed. Each integer is the
        dilation performed at that convolution. Typically, this list follows
        the pattern of first increasing in dilation rate, and then
        decreasing again. Default is `(1, 2, 4, 8, 4, 2, 1)`.
    - `use_norm`: Whether to use group normalization. Default is `True`.
    - `num_groups`: The number of groups to use for group normalization.
        Default is `1`.
    - `use_bias`: Whether to use bias in the convolutional layers. Default is
        `True`.
    - `zero_bias_init`: Whether to initialise the bias to zero. Default is
        `False`.
    """
    self.kernel_size = kernel_size
    self.dilation_rates = dilation_rates
    self.use_norm = use_norm
    self.num_groups = num_groups
    self.use_bias = use_bias
    self.zero_bias_init = zero_bias_init
__call__ ¤
__call__(
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal[
        "periodic", "dirichlet", "neumann"
    ],
    key: PRNGKeyArray
) -> DilatedResBlock
Source code in pdequinox/blocks/_dilated_res_block.py 
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
def __call__(
    self,
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal["periodic", "dirichlet", "neumann"],
    key: PRNGKeyArray,
) -> DilatedResBlock:
    return DilatedResBlock(
        num_spatial_dims=num_spatial_dims,
        in_channels=in_channels,
        out_channels=out_channels,
        activation=activation,
        kernel_size=self.kernel_size,
        dilation_rates=self.dilation_rates,
        boundary_mode=boundary_mode,
        key=key,
        use_norm=self.use_norm,
        num_groups=self.num_groups,
        use_bias=self.use_bias,
        zero_bias_init=self.zero_bias_init,
    )

pdequinox.blocks.LinearConvBlockFactory ¤

Bases: BlockFactory

Source code in pdequinox/blocks/_linear_conv_block.py 
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
class LinearConvBlockFactory(BlockFactory):
    kernel_size: int
    use_bias: bool

    def __init__(
        self,
        *,
        kernel_size: int = 3,
        use_bias: bool = True,
    ):
        """
        Factory for creating `LinearConvBlock` instances.

        **Arguments:**

        - `kernel_size`: The size of the convolutional kernel. Default is `3`.
        - `use_bias`: Whether to use bias in the convolutional layers. Default is
            `True`.
        """
        self.kernel_size = kernel_size
        self.use_bias = use_bias

    def __call__(
        self,
        num_spatial_dims: int,
        in_channels: int,
        out_channels: int,
        *,
        activation: Callable,  # unused
        boundary_mode: Literal["periodic", "dirichlet", "neumann"],
        key: PRNGKeyArray,
    ):
        return LinearConvBlock(
            num_spatial_dims=num_spatial_dims,
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=self.kernel_size,
            boundary_mode=boundary_mode,
            use_bias=self.use_bias,
            key=key,
        )
__init__ ¤
__init__(*, kernel_size: int = 3, use_bias: bool = True)

Factory for creating LinearConvBlock instances.

Arguments:

  • kernel_size: The size of the convolutional kernel. Default is 3.
  • use_bias: Whether to use bias in the convolutional layers. Default is True.
Source code in pdequinox/blocks/_linear_conv_block.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def __init__(
    self,
    *,
    kernel_size: int = 3,
    use_bias: bool = True,
):
    """
    Factory for creating `LinearConvBlock` instances.

    **Arguments:**

    - `kernel_size`: The size of the convolutional kernel. Default is `3`.
    - `use_bias`: Whether to use bias in the convolutional layers. Default is
        `True`.
    """
    self.kernel_size = kernel_size
    self.use_bias = use_bias
__call__ ¤
__call__(
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal[
        "periodic", "dirichlet", "neumann"
    ],
    key: PRNGKeyArray
)
Source code in pdequinox/blocks/_linear_conv_block.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def __call__(
    self,
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,  # unused
    boundary_mode: Literal["periodic", "dirichlet", "neumann"],
    key: PRNGKeyArray,
):
    return LinearConvBlock(
        num_spatial_dims=num_spatial_dims,
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=self.kernel_size,
        boundary_mode=boundary_mode,
        use_bias=self.use_bias,
        key=key,
    )

pdequinox.blocks.LinearConvDownBlockFactory ¤

Bases: BlockFactory

Source code in pdequinox/blocks/_linear_conv_down_block.py 
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
class LinearConvDownBlockFactory(BlockFactory):
    kernel_size: int
    factor: int
    use_bias: bool

    def __init__(
        self,
        *,
        kernel_size: int = 3,
        factor: int = 2,
        use_bias: bool = True,
    ):
        """
        Factory for creating `LinearConvDownBlock` instances.

        **Arguments:**

        - `kernel_size`: The size of the convolutional kernel. Default is `3`.
        - `factor`: The downsampling factor. Default is `2`. This will become
            the stride of the convolution.
        - `use_bias`: Whether to use bias after the convolution. Default
            is `True`.
        """
        self.kernel_size = kernel_size
        self.factor = factor
        self.use_bias = use_bias

    def __call__(
        self,
        num_spatial_dims: int,
        in_channels: int,
        out_channels: int,
        *,
        activation: Callable,  # unused
        boundary_mode: Literal["periodic", "dirichlet", "neumann"],
        key: PRNGKeyArray,
    ):
        return LinearConvDownBlock(
            num_spatial_dims=num_spatial_dims,
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=self.kernel_size,
            factor=self.factor,
            boundary_mode=boundary_mode,
            use_bias=self.use_bias,
            key=key,
        )
__init__ ¤
__init__(
    *,
    kernel_size: int = 3,
    factor: int = 2,
    use_bias: bool = True
)

Factory for creating LinearConvDownBlock instances.

Arguments:

  • kernel_size: The size of the convolutional kernel. Default is 3.
  • factor: The downsampling factor. Default is 2. This will become the stride of the convolution.
  • use_bias: Whether to use bias after the convolution. Default is True.
Source code in pdequinox/blocks/_linear_conv_down_block.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def __init__(
    self,
    *,
    kernel_size: int = 3,
    factor: int = 2,
    use_bias: bool = True,
):
    """
    Factory for creating `LinearConvDownBlock` instances.

    **Arguments:**

    - `kernel_size`: The size of the convolutional kernel. Default is `3`.
    - `factor`: The downsampling factor. Default is `2`. This will become
        the stride of the convolution.
    - `use_bias`: Whether to use bias after the convolution. Default
        is `True`.
    """
    self.kernel_size = kernel_size
    self.factor = factor
    self.use_bias = use_bias
__call__ ¤
__call__(
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal[
        "periodic", "dirichlet", "neumann"
    ],
    key: PRNGKeyArray
)
Source code in pdequinox/blocks/_linear_conv_down_block.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def __call__(
    self,
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,  # unused
    boundary_mode: Literal["periodic", "dirichlet", "neumann"],
    key: PRNGKeyArray,
):
    return LinearConvDownBlock(
        num_spatial_dims=num_spatial_dims,
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=self.kernel_size,
        factor=self.factor,
        boundary_mode=boundary_mode,
        use_bias=self.use_bias,
        key=key,
    )

pdequinox.blocks.LinearConvUpBlockFactory ¤

Bases: BlockFactory

Source code in pdequinox/blocks/_linear_conv_up_block.py 
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
class LinearConvUpBlockFactory(BlockFactory):
    kernel_size: int
    factor: int
    use_bias: bool
    output_padding: int

    def __init__(
        self,
        *,
        kernel_size: int = 3,
        factor: int = 2,
        use_bias: bool = True,
        output_padding: int = 1,
    ):
        """
        Factory for creating `LinearConvUpBlock` instances.

        **Arguments:**

        - `kernel_size`: The size of the convolutional kernel. Default is `3`.
        - `factor`: The upsampling factor. Default is `2`. This will become
            the stride of the transposed convolution. Set this to the same value
            as in a corresponding `LinearConvDownBlockFactory` instance.
        - `use_bias`: Whether to use bias after the convolution. Default
            is `True`.
        - `output_padding`: The amount of additional padding used by the
            transposed convolution. Use this to resolve the ambiguity that the
            result of an integer division with `factor` is not bijective. If you
            have `factor=2` and work with spatial dimensions divisible by `2`,
            set this to `1`. Default is `1`.
        """
        self.kernel_size = kernel_size
        self.factor = factor
        self.use_bias = use_bias
        self.output_padding = output_padding

    def __call__(
        self,
        num_spatial_dims: int,
        in_channels: int,
        out_channels: int,
        *,
        activation: Callable,  # unused
        boundary_mode: Literal["periodic", "dirichlet", "neumann"],
        key: PRNGKeyArray,
    ):
        return LinearConvUpBlock(
            num_spatial_dims=num_spatial_dims,
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=self.kernel_size,
            factor=self.factor,
            output_padding=self.output_padding,
            boundary_mode=boundary_mode,
            use_bias=self.use_bias,
            key=key,
        )
__init__ ¤
__init__(
    *,
    kernel_size: int = 3,
    factor: int = 2,
    use_bias: bool = True,
    output_padding: int = 1
)

Factory for creating LinearConvUpBlock instances.

Arguments:

  • kernel_size: The size of the convolutional kernel. Default is 3.
  • factor: The upsampling factor. Default is 2. This will become the stride of the transposed convolution. Set this to the same value as in a corresponding LinearConvDownBlockFactory instance.
  • use_bias: Whether to use bias after the convolution. Default is True.
  • output_padding: The amount of additional padding used by the transposed convolution. Use this to resolve the ambiguity that the result of an integer division with factor is not bijective. If you have factor=2 and work with spatial dimensions divisible by 2, set this to 1. Default is 1.
Source code in pdequinox/blocks/_linear_conv_up_block.py 
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
def __init__(
    self,
    *,
    kernel_size: int = 3,
    factor: int = 2,
    use_bias: bool = True,
    output_padding: int = 1,
):
    """
    Factory for creating `LinearConvUpBlock` instances.

    **Arguments:**

    - `kernel_size`: The size of the convolutional kernel. Default is `3`.
    - `factor`: The upsampling factor. Default is `2`. This will become
        the stride of the transposed convolution. Set this to the same value
        as in a corresponding `LinearConvDownBlockFactory` instance.
    - `use_bias`: Whether to use bias after the convolution. Default
        is `True`.
    - `output_padding`: The amount of additional padding used by the
        transposed convolution. Use this to resolve the ambiguity that the
        result of an integer division with `factor` is not bijective. If you
        have `factor=2` and work with spatial dimensions divisible by `2`,
        set this to `1`. Default is `1`.
    """
    self.kernel_size = kernel_size
    self.factor = factor
    self.use_bias = use_bias
    self.output_padding = output_padding
__call__ ¤
__call__(
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal[
        "periodic", "dirichlet", "neumann"
    ],
    key: PRNGKeyArray
)
Source code in pdequinox/blocks/_linear_conv_up_block.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def __call__(
    self,
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,  # unused
    boundary_mode: Literal["periodic", "dirichlet", "neumann"],
    key: PRNGKeyArray,
):
    return LinearConvUpBlock(
        num_spatial_dims=num_spatial_dims,
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=self.kernel_size,
        factor=self.factor,
        output_padding=self.output_padding,
        boundary_mode=boundary_mode,
        use_bias=self.use_bias,
        key=key,
    )

pdequinox.blocks.BlockFactory ¤

Bases: Module, ABC

Source code in pdequinox/blocks/_base_block.py 
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
class BlockFactory(eqx.Module, ABC):
    @abstractmethod
    def __call__(
        self,
        num_spatial_dims: int,
        in_channels: int,
        out_channels: int,
        *,
        activation: Callable,
        boundary_mode: Literal["periodic", "dirichlet", "neumann"],
        key: PRNGKeyArray,
    ) -> Block:
        """
        Construct a block (= `equinox.Module`)

        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions. For example
            traditional, convolutions for image processing have this set to `2`.
        - `in_channels`: The number of input channels.
        - `out_channels`: The number of output channels.
        - `activation`: The activation function to use. For example
            `jax.nn.relu`.
        - `boundary_mode`: The boundary mode to use. For example `"periodic"`.
            (Keyword only argument.)
        - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
            initialisation. (Keyword only argument.)
        - ` `: Additional keyword arguments to pass to the boundary
            mode constructor.
        """
        pass
__call__ abstractmethod ¤
__call__(
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal[
        "periodic", "dirichlet", "neumann"
    ],
    key: PRNGKeyArray
) -> Block

Construct a block (= equinox.Module)

Arguments:

  • num_spatial_dims: The number of spatial dimensions. For example traditional, convolutions for image processing have this set to 2.
  • in_channels: The number of input channels.
  • out_channels: The number of output channels.
  • activation: The activation function to use. For example jax.nn.relu.
  • boundary_mode: The boundary mode to use. For example "periodic". (Keyword only argument.)
  • key: A jax.random.PRNGKey used to provide randomness for parameter initialisation. (Keyword only argument.)
  • : Additional keyword arguments to pass to the boundary mode constructor.
Source code in pdequinox/blocks/_base_block.py 
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
@abstractmethod
def __call__(
    self,
    num_spatial_dims: int,
    in_channels: int,
    out_channels: int,
    *,
    activation: Callable,
    boundary_mode: Literal["periodic", "dirichlet", "neumann"],
    key: PRNGKeyArray,
) -> Block:
    """
    Construct a block (= `equinox.Module`)

    **Arguments:**

    - `num_spatial_dims`: The number of spatial dimensions. For example
        traditional, convolutions for image processing have this set to `2`.
    - `in_channels`: The number of input channels.
    - `out_channels`: The number of output channels.
    - `activation`: The activation function to use. For example
        `jax.nn.relu`.
    - `boundary_mode`: The boundary mode to use. For example `"periodic"`.
        (Keyword only argument.)
    - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
        initialisation. (Keyword only argument.)
    - ` `: Additional keyword arguments to pass to the boundary
        mode constructor.
    """
    pass