pyhgf.model.conv.im2col_adapter#

pyhgf.model.conv.im2col_adapter(filter_shape, strides=(1, 1), padding='SAME')[source]#

Build a frozen patch extractor: images in, one row per patch out.

Forward, x of shape (batch, C, H, W) becomes (batch, n_patches, C*kh*kw) via jax.lax.conv_general_dilated_patches, the same “one row per sample” layout DeepNetworkAdapter expects for token positions, so the following per-patch network needs no special-casing. Backward (fold / col2im) is the geometric adjoint: each of the kh*kw kernel taps is a strided slice of the (padded) image, so its gradient is scattered back with jax.numpy.ndarray.at.add() at the same stride, a closed-form sum over overlapping patches unrolled over the statically-sized kernel taps. The adjoint is written out rather than obtained by differentiating the forward at call time; it matches jax.vjp of this forward (tests/test_conv.py).

Parameters:

filter_shape (tuple[int, int])
strides (tuple[int, int])
padding (str | tuple[tuple[int, int], tuple[int, int]])

Return type:

EquinoxAdapter