search

patrick-kidger / equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
Apache License 2.0
2.13k stars 143 forks source link

Comparable PixelShuffle functionality? #812

Closed mathDR closed 3 months ago

mathDR commented 3 months ago

Is there a way to add a bespoke layer? In particular, I am trying to reproduce a torch model that has a PixelShuffle(2) layer embedded in a nn.Sequential() call.

What would be the best way to do that?

lockwo commented 3 months ago

You can write a module that has the functionality you want, and use it in your bigger model. Like https://github.com/patrick-kidger/equinox/issues/809#issuecomment-2303854615 for example

mathDR commented 3 months ago

Awesome! Thanks for the example, that is totally clear,