Model parameters end up as constants in the IR when exporting to stable-hlo

[navdeepkk-polymagelabs]

I tried exporting ResNet from Google Scenic (https://github.com/google-research/scenic) library to stable hlo. This lowering captured all the model parameters as constants. Is there a way to capture the model parameters as the function arguments rather than constants?

I am using the following jit call to lower a given model to stable-hlo:

mlir_stable_hlo = jax.jit(model).lower(*input_signature).as_text()

Thanks!