code: def optimize(
      self,
      pretrained: nnx.Linear,
      x: jax.Array,
      *,
      num_steps: int = 100,
      debug: bool = False,
  ):
      q_hparams, rest, static = self.split(QHParam, ...)
      tx = optax.adam(1e-3)
      opt_state = tx.init(q_hparams)
  
      print(jax.tree_map(lambda x: x.shape, q_hparams))
  
      @jax.jit
      def optimization_step(
          q_hparams: nnx.State,
          rest: nnx.State,
          opt_state: optax.OptState,
          x: jax.Array,
      ):
          print("JITTING")
  
          def loss_fn(q_hparams: nnx.State):
              model = static.merge(q_hparams, rest)
              model.qkernel = model.quantize(
                  pretrained.kernel, 8, jnp.uint8
              )
              assert pretrained.bias is not None
              model.qbias = model.quantize(
                  pretrained.bias, 16, jnp.uint16
              )
  
              y_quant = model(x)
              y_unquant = pretrained(x)
              loss = jnp.mean((y_unquant - y_quant) ** 2)
              return loss
  
          loss, grads = jax.value_and_grad(loss_fn)(q_hparams)
  
          updates, opt_state = tx.update(grads, opt_state, q_hparams)
          q_hparams = optax.apply_updates(q_hparams, updates)  # type: ignore
  
          return q_hparams, opt_state, loss
  
      for step in range(num_steps):
          q_hparams, opt_state, loss = optimization_step(
              q_hparams, rest, opt_state, x
          )
          if debug and step % (num_steps / 10) == 0:
              print(f"Step {step} - loss: {loss}")
  
      self.update(q_hparams)
  
      self.qkernel = self.quantize(pretrained.kernel, 8, jnp.uint8)
      assert pretrained.bias is not None
      self.qbias = self.quantize(pretrained.bias, 16, jnp.uint16)
  
imports: import typing
  from functools import partial
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  import optax
  from datasets import load_dataset
  from flax.experimental import nnx
label_desc: Use image node_id 5 for calling the QLinear obj's underlying member method code with example usage: obj.optimize(pretrained, x) and returns: q_hparams,  opt_state,  loss, loss