message: avoid python circular import
diff: import torch
  -import intel_pytorch_extension as ipex
  import _torch_ipex as core
  from torch.jit._recursive import wrap_cpp_module
  @@ -14,26 +13,27 @@ def script_(obj, optimize=None, _frames_up=0, _rcb=None):
  jit_m = orig_script(obj, optimize=optimize, _frames_up=_frames_up+1, _rcb=_rcb)
  torch.jit.script = script_
  - if core.get_jit_opt() and hasattr(jit_m, '_c'):
  + mix_state = core.get_mix_bf16_fp32();
  # Disable mix precision in model fusion, since mixed precision cannot
  # bring any benefits for inference, but will lead to loss of accuracy
  - orig_mixed_type = ipex.get_auto_mix_precision()
  - ipex.enable_auto_mix_precision(None)
  + core.disable_mix_bf16_fp32()
  + if core.get_jit_opt() and hasattr(jit_m, '_c'):
  jit_m = wrap_cpp_module(torch._C._jit_pass_fold_convbn(jit_m._c))
  - ipex.enable_auto_mix_precision(orig_mixed_type)
  + if mix_state:
  + core.enable_mix_bf16_fp32()
  return jit_m
  def trace_(func, example_inputs, *args, **kwargs):
  # Disable mix precision. torch.jit.trace will check the traced output
  # against what is expected. Since mix precision will lead to
  # loss of accuracy, this will raise warning during torch.jit.trace
  - orig_mixed_type = ipex.get_auto_mix_precision()
  - ipex.enable_auto_mix_precision(None)
  + mix_state = core.get_mix_bf16_fp32()
  + core.disable_mix_bf16_fp32()
  jit_m = orig_trace(func, example_inputs, *args, **kwargs)
  -
  if core.get_jit_opt() and hasattr(jit_m, '_c'):
  jit_m = wrap_cpp_module(torch._C._jit_pass_fold_convbn(jit_m._c))
  - ipex.enable_auto_mix_precision(orig_mixed_type)
  + if mix_state:
  + core.enable_mix_bf16_fp32()
  return jit_m
  