Is there a function like torch.nn.fold in JAX?

[ZhongkuiMa]

Hi Everyone,

I have known that conv_general_dilated_patches can realize a function like the torch.nn.unfold.
But I need an implementation like the torch.nn.fold.
Is there a function in JAX that can achieve this?
Or should I implement it by myself? : )

Thank you