code: def quantization(x, s, z, alpha_q, beta_q):
      x_q = jnp.round(1 / s * x + z, decimals=0)
      x_q = jnp.clip(x_q, a_min=alpha_q, a_max=beta_q)
  
      return x_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 12 for calling a global function with example usage: quantization(x, s, z, alpha_q, beta_q) and returns: x_q