Custom VJP with class / method

[karan-dalal]

I want to implement a custom VJP for a class, which implements a nn.module (using Flax). An example:

class FlaxTTTModule(nn.Module):
    config: Any = None
    dtype: jnp.dtype = jnp.float32

    def setup(self):        
        self.InnerNet = FlaxTTTInnerNet(self.config, self.dtype)
        variable = self.InnerNet.init(jax.random.PRNGKey(0), jnp.ones([1, self.config.n_embd // self.config.n_head]), train=False)
        inner_params = variable['params']

    def __call__(self, batch, train=False):
        replicate_inner_params = self.replicate_param(self.inner_params, batch.shape[0])
        return flax_ttt_default_forward(self.InnerNet.apply, self.config, batch, replicate_inner_params, train)

Outside of the class, I've defined my custom VJP:

@partial(custom_vjp, nondiff_argnums=(0, 1))
def flax_ttt_default_forward(InnerNet_apply, config, batch, replicate_inner_params, train=False):
          inner_input_reshape = inner_input.reshape(config.n_head, -1)
          inner_input_multi_head_transformed = InnerNet_apply(params, inner_input_reshape, train=False)

However, I get this error:

TypeError: Value <bound method FlaxTTTModule.apply_InnerNet of FlaxTTTModule(
    # attributes
    config = GPTConfig(block_size=32, vocab_size=50304, n_layer=1, n_head=1, n_embd=8, dropout=0.0, bias=False, inner_net='linear', train_init=True, sgd=True, ssl_type='AB', n_btn=12)
    dtype = float32
    # children
    InnerNet = FlaxTTTInnerNet(
        # attributes
        config = GPTConfig(block_size=32, vocab_size=50304, n_layer=1, n_head=1, n_embd=8, dropout=0.0, bias=False, inner_net='linear', train_init=True, sgd=True, ssl_type='AB', n_btn=12)
        dtype = float32
    )
)> with type <class 'method'> is not a valid JAX type

I assume this is because I cannot pass a method into the custom VJP, even with defining nondiff_arguments? Is there a way to solve this - should I be moving my custom VJP inside the class?

Any tips helpful, thanks!

[jakevdp]

It's hard to piece together exactly what code you're executing when you see this error. Would it be possible for you to put together a minimal reproducible example?

[karan-dalal]

Hey - sorry about that. I've realized what my error is: I'm trying to return a non-JITable type from my custom VJP forward pass (the one used during backward) as an intermediate / residual.

I think I should be able to solve this by converting my method into a partial? here's an example of what I'm trying to do:

class FlaxModule contains an InnerNetwork, which has an .apply method. In the custom VJP, I'd like to apply pass in the InnerNetwork.apply method and use it. However, I'd also like to then pass this method as an intermediate to be utilized by the backward pass.

I'm also wondering if I should define my VJP inside or outside my class?

[jakevdp]

I generally think it's cleaner to define custom_jvp on simple functions outside the class definition, and then have class methods call out to those utility functions where appropriate.

[karan-dalal]

Got it. I got what I was looking for, just need to pass in non-diff arguments separately from the residuals to the backward pass. Thanks!