docstr: return quantized x by pytorch
docstr_tok: ["return","quantized","x","by","pytorch"]
code: def torch_quantize(
      x: np.ndarray,
      scale: np.ndarray,
      zero_point: np.ndarray,
      axis: int = None,
      quant_dtype: np.dtype = None,
  ) -> torch.Tensor:
      """
      return quantized x by pytorch
      """
  
      # quantization data type is either inferred from `zero_point`,
      # or explicitly provided
      if zero_point is not None:
          quant_dtype = zero_point.dtype
      assert quant_dtype is not None
  
      # if scale is scalar, then axis must be None
      # if scale is not scalar, then axis must have a value
      assert (len(scale.shape) == 0) != (axis is not None)
  
      x_torch = torch.from_numpy(x).to(torch.float32)
      s_torch = torch.from_numpy(scale).to(torch.float32)
      zp_torch = (
          torch.zeros(scale.shape, dtype=torch.int)
          if zero_point is None
          else torch.from_numpy(zero_point)
      )
      dtype_torch = (
          torch.quint8 if quant_dtype == np.uint8 else torch.qint8
      )
  
      output: np.ndarray
      if axis is None:
          output = torch.quantize_per_tensor(
              x_torch, s_torch, zp_torch, dtype_torch
          )
      else:
          if axis < 0:
              axis += len(x.shape)
          output = torch.quantize_per_channel(
              x_torch, s_torch, zp_torch, axis, dtype_torch
          )
      return output
  
imports: import itertools
  from typing import Tuple
  import numpy
  import pytest
  from coremltools._deps import _HAS_TORCH, MSG_TORCH_NOT_FOUND
  from coremltools.converters.mil.mil import Builder
  from coremltools.converters.mil.mil.ops.tests.iOS17 import backends
  from coremltools.converters.mil.mil.ops.tests.testing_utils import run_compare_builder
  from coremltools.converters.mil.mil.types import builtin_to_string, numpy_type_to_builtin_type
  from coremltools.converters.mil.testing_reqs import BackendConfig, compute_units
  from coremltools.converters.mil.testing_utils import ssa_fn
label_desc: Use image node_id 2 for calling the TestQuantizationBase obj's underlying member method code with example usage: obj.torch_quantize(x, scale, zero_point, axis, quant_dtype) and returns: output