code: def quantization_matrix_multiplication_int8(
      X_q, W_q, b_q, s_X, z_X, s_W, z_W, s_b, z_b, s_Y, z_Y
  ):
      p = W_q.shape[0]
  
      # Y_q_simulated is FP32
      Y_q_simulated = (
          z_Y
          + (s_b / s_Y * (b_q.astype(jnp.int32) - z_b))
          + (
              (s_X * s_W / s_Y)
              * (
                  jnp.matmul(
                      X_q.astype(jnp.int32), W_q.astype(jnp.int32)
                  )
                  - z_W
                  * jnp.sum(
                      X_q.astype(jnp.int32), axis=1, keepdims=True
                  )
                  - z_X
                  * jnp.sum(
                      W_q.astype(jnp.int32), axis=0, keepdims=True
                  )
                  + p * z_X * z_W
              )
          )
      )
  
      Y_q_simulated = jnp.round(Y_q_simulated, decimals=0)
      Y_q_simulated = jnp.clip(Y_q_simulated, a_min=-128, a_max=127)
      Y_q_simulated = Y_q_simulated.astype(jnp.int8)
  
      return Y_q_simulated
  
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 17 for calling a global function with example usage: quantization_matrix_multiplication_int8(X_q, W_q, b_q, s_X, z_X, s_W, z_W, s_b, z_b, s_Y, z_Y) and returns: Y_q_simulated