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!

[cgarciae]

Hey @navdeepkk-polymagelabs, you will have to give a bit more context. Its not clear what model is here in this case. Usually you create a "forward" function or similar and lower that:

def forward(params, x):
  return model.apply({'params': params, x)
  
mlir_stable_hlo = jax.jit(forward).lower(params, x).as_text()

[navdeepkk-polymagelabs]

Thanks @cgarciae. That was pretty much it. I wasn't JIT'ing the the forward function properly. I had

params = ...

def forward(x):
  return model.apply({'params': params, x)
  
mlir_stable_hlo = jax.jit(forward).lower(x).as_text()

Doing this the params ended up as constants in the IR.