Inner workings of jit(model.apply)

[pyatsysh]

Can somebody please explain to me

1. why jit(model.apply) throws a concretisation error if nn.Module fields are pushed through any jnp functions. And at the same time, same operations with np functions don't raise jit errors?

2. Is there a workaround to make jit(model.apply) work when pushing model fields through jnp functions?

A minimal example follows. Thank you!!

import flax.linen as nn
import jax.numpy as jnp
import numpy as np
from jax import random, jit
key = random.PRNGKey(0)


class demo(nn.Module):
    """
    jax concretisation error example
    """
    N: int = 5

    @nn.compact
    def __call__(self, x):

        # jit(model.apply)(x) works:
        # N = np.array(self.N, dtype=int)

        # jit(model.apply)(x) breaks:
        N = jnp.array(self.N, dtype = int)

        q = self.param('q', 
                       nn.initializers.constant(1.), 
                       (N,))
        
        return x*q.sum()



x = random.uniform(key, (2,3,4))
model = demo(5)
params = model.init(key, x)
_  = model.apply(params, x) # always passes
_ = jit(model.apply)(params,x) # fails with concretisation error