Rematerialize all dots except dots in a specific submodule

[hr0nix]

Hi!

Let's say my model has the following structure:

block 1:
  subblock 1
  subblock 2
  subblock_no_remat
...
block n:
  subblock 1
  subblock 2
  subblock_no_remat

I need to setup the following rematerialization scheme: save inputs and outputs of all dots, but only in subblock_1 and subblock_2, not subblock_no_remat. Is there a way to achieve this in jax?

I was hoping that wrapping the whole model in remat(policy=dots_saveable) and subblock_no_remat in remat(policy=nothing_saveable) will do the job, but from the looks of it that's not how rematerialization nesting works (actually, I can't figure out exactly what it does, but I can see that internals of subblock_no_remat are being saved for each layer).

Is there a way to achieve what I need?