How to get image pyramid in jax?

[wztdream]

Hi,
I have a project which need convert an input image to a pyramid, then extract features from each scale in the pyramid using CNN. I'm using flax by the way.

In below code, I have tried pyramid_gaussian from skimage which can not handle a batch of images, and it is non-jit-able, which will trigger TracerArrayConversionError if we try to jit a function that calls this module.

So how to get image pyramid in jax?

Or maybe we can treat image pyramid as a data processing step, and use tensorflow or pytorch pyramid funtion to get a batch of image pyrmaid, and take the pyramid as input directly to the module, what do you think?

from skimage.transform import pyramid_gaussian
class Foo(nn.Module):
    @nn.compact
    def __call__(self, image):
        pyra = tuple(pyramid_gaussian(image, downscale=2)) # here we use non-jit-able funcion
        pyra = [scale[None, ...] for scale in pyra]
        r = 0.0
        for scale in pyra:
            x = nn.Conv(2, kernel_size=(1, 1), strides=(1, 1), padding="SAME")(scale)
            r += x.mean(axis=(1, 2, 3))
        return r