code: def optimize2(
      self,
      pretrained: nnx.Linear,
      X: jax.Array,
  ):
      W = pretrained.kernel
      b = pretrained.bias
      assert b is not None
  
      # X
      alpha_X = jnp.min(X)
      beta_X = jnp.max(X)
      s_X, z_X = generate_quantization_int8_constants(
          alpha=alpha_X, beta=beta_X
      )
      X_q = quantization_int8(x=X, s=s_X, z=z_X)
      X_q_dq = dequantization(x_q=X_q, s=s_X, z=z_X)
  
      # W
      alpha_W = jnp.min(W)
      beta_W = jnp.max(W)
      s_W, z_W = generate_quantization_int8_constants(
          alpha=alpha_W, beta=beta_W
      )
      W_q = quantization_int8(x=W, s=s_W, z=z_W)
      W_q_dq = dequantization(x_q=W_q, s=s_W, z=z_W)
  
      # b
      alpha_b = jnp.min(b)
      beta_b = jnp.max(b)
      s_b, z_b = generate_quantization_int8_constants(
          alpha=alpha_b, beta=beta_b
      )
      b_q = quantization_int8(x=b, s=s_b, z=z_b)
      b_q_dq = dequantization(x_q=b_q, s=s_b, z=z_b)
  
      # Y
      Y = jnp.matmul(X, W) + b
      alpha_Y = jnp.min(Y)
      beta_Y = jnp.max(Y)
      s_Y, z_Y = generate_quantization_int8_constants(
          alpha=alpha_Y, beta=beta_Y
      )
      Y_q = quantization_int8(x=Y, s=s_Y, z=z_Y)
  
      Y_prime = jnp.matmul(X_q_dq, W_q_dq) + b_q_dq
      Y_prime_q = quantization_int8(x=Y_prime, s=s_Y, z=z_Y)
      Y_prime_q_dq = dequantization(x_q=Y_prime_q, s=s_Y, z=z_Y)
  
      print("Expected FP32 Y:")
      print(Y)
      print("Expected FP32 Y Quantized:")
      print(Y_q)
  
      Y_q_simulated = quantization_matrix_multiplication_int8(
          X_q=X_q,
          W_q=W_q,
          b_q=b_q,
          s_X=s_X,
          z_X=z_X,
          s_W=s_W,
          z_W=z_W,
          s_b=s_b,
          z_b=z_b,
          s_Y=s_Y,
          z_Y=z_Y,
      )
      Y_simulated = dequantization(x_q=Y_q_simulated, s=s_Y, z=z_Y)
  
      print(
          "Expected Quantized Y_q from Quantized Matrix Multiplication:"
      )
      print(Y_q_simulated)
      print(
          "Expected Quantized Y_q from Quantized Matrix Multiplication Dequantized:"
      )
      print(Y_simulated)
  
      # Ensure the algorithm implementation is correct
      assert jnp.array_equal(Y_simulated, Y_prime_q_dq)
      assert jnp.array_equal(Y_q_simulated, Y_prime_q)
  
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 11 for calling a global function with example usage: optimize2(pretrained, X) without return types