content: # encoding: utf-8
  import logging
  import operator
  import warnings
  
  import torch  # isort:skip
  from typing import cast, Iterable, List, Optional, Sequence
  
  import torch.nn as nn
  from torch.fx.passes.shape_prop import _extract_tensor_metadata, TensorMetadata
  
  from . import acc_utils
  from .acc_normalizer import (
      register_acc_op,
      register_acc_op_mapping,
      register_custom_acc_mapper_fn,
  )
  from .acc_op_properties import AccOpProperty, register_acc_op_properties
  
  logger: logging.Logger = logging.getLogger(__name__)
  
  this_arg_is_optional = True
  move_to_qparams = True
  dont_move_to_qparams = False
  
  # A proxy embedding size. We use this for tracing proxy operators using XL
  # weights which we can't load into memory (because they're too large), we
  # instead substitute a smaller weight with embedding size =
  # PROXY_EMBEDDING_SIZE.
  PROXY_EMBEDDING_SIZE = 8
  
  
  @register_acc_op_mapping(op_and_target=("call_function", nn.functional.linear))
  @register_acc_op
  def linear(*, input, weight, bias):
      return nn.functional.linear(input=input, weight=weight, bias=bias)
  
  
  @register_acc_op_properties(AccOpProperty.quantized)
  @register_acc_op
  def quantized_linear(*, input, weight, bias, acc_out_ty=None):
      assert acc_out_ty is not None
      qparams = acc_out_ty.qparams
      return nn.quantized.functional.linear(
          input,
          weight,
          bias,
          qparams["scale"],
          qparams["zero_point"],
      )
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op_mapping(
      op_and_target=("call_method", "flatten"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("start_dim", "start_dim", this_arg_is_optional),
          ("end_dim", "end_dim", this_arg_is_optional),
      ],
  )
  @register_acc_op_mapping(op_and_target=("call_function", torch.flatten))
  @register_acc_op
  def flatten(*, input, start_dim=0, end_dim=-1):
      return torch.flatten(input=input, start_dim=start_dim, end_dim=end_dim)
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op_mapping(
      op_and_target=("call_method", "squeeze"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("dim", "dim", this_arg_is_optional),
      ],
  )
  @register_acc_op_mapping(
      op_and_target=("call_function", torch.squeeze),
      arg_replacement_tuples=[
          ("input", "input"),
          ("dim", "dim", this_arg_is_optional),
      ],
  )
  @register_acc_op
  def squeeze(*, input, dim=None):
      if dim is None:
          return input.squeeze()
      return input.squeeze(dim=dim)
  
  
  @register_acc_op_mapping(op_and_target=("call_function", nn.functional.embedding))
  @register_acc_op
  def embedding(
      *,
      input,
      weight,
      padding_idx,
      max_norm,
      norm_type,
      scale_grad_by_freq,
      sparse,
  ):
      return torch.nn.functional.embedding(**locals())
  
  
  @register_acc_op_mapping(op_and_target=("call_function", nn.functional.max_pool1d))
  @register_acc_op
  def max_pool1d(
      *,
      input,
      kernel_size,
      stride,
      padding,
      dilation,
      ceil_mode,
      return_indices,
  ):
      return nn.functional.max_pool1d(
          input=input,
          kernel_size=kernel_size,
          stride=stride,
          padding=padding,
          dilation=dilation,
          ceil_mode=ceil_mode,
          return_indices=return_indices,
      )
  
  
  @register_acc_op_mapping(op_and_target=("call_function", nn.functional.max_pool2d))
  @register_acc_op
  def max_pool2d(
      *,
      input,
      kernel_size,
      stride,
      padding,
      dilation,
      ceil_mode,
      return_indices,
  ):
      return nn.functional.max_pool2d(
          input=input,
          kernel_size=kernel_size,
          stride=stride,
          padding=padding,
          dilation=dilation,
          ceil_mode=ceil_mode,
          return_indices=return_indices,
      )
  
  
  @register_acc_op_mapping(op_and_target=("call_function", nn.functional.max_pool3d))
  @register_acc_op
  def max_pool3d(
      *, input, kernel_size, stride, padding, dilation, ceil_mode, return_indices
  ):
      return nn.functional.max_pool3d(
          input=input,
          kernel_size=kernel_size,
          stride=stride,
          padding=padding,
          dilation=dilation,
          ceil_mode=ceil_mode,
          return_indices=return_indices,
      )
  
  
  @register_acc_op_mapping(op_and_target=("call_function", nn.functional.normalize))
  @register_acc_op
  def normalize(*, input, p, dim, eps, out):
      return nn.functional.normalize(
          input=input,
          p=p,
          dim=dim,
          eps=eps,
          out=out,
      )
  
  
  @register_acc_op_mapping(
      op_and_target=("call_function", nn.functional.adaptive_avg_pool2d)
  )
  @register_acc_op
  def adaptive_avg_pool2d(*, input, output_size):
      return nn.functional.adaptive_avg_pool2d(input=input, output_size=output_size)
  
  
  @register_acc_op_mapping(
      op_and_target=("call_function", nn.functional.adaptive_avg_pool3d)
  )
  @register_acc_op
  def adaptive_avg_pool3d(*, input, output_size):
      return nn.functional.adaptive_avg_pool3d(input=input, output_size=output_size)
  
  
  @register_acc_op_mapping(op_and_target=("call_function", nn.functional.avg_pool1d))
  @register_acc_op
  def avg_pool1d(
      *,
      input,
      kernel_size,
      stride,
      padding,
      ceil_mode,
      count_include_pad,
  ):
      return nn.functional.avg_pool1d(
          input=input,
          kernel_size=kernel_size,
          stride=stride,
          padding=padding,
          ceil_mode=ceil_mode,
          count_include_pad=count_include_pad,
      )
  
  
  @register_acc_op_mapping(op_and_target=("call_function", nn.functional.avg_pool2d))
  @register_acc_op
  def avg_pool2d(
      *,
      input,
      kernel_size,
      stride,
      padding,
      ceil_mode,
      count_include_pad,
      divisor_override,
  ):
      return nn.functional.avg_pool2d(
          input=input,
          kernel_size=kernel_size,
          stride=stride,
          padding=padding,
          ceil_mode=ceil_mode,
          count_include_pad=count_include_pad,
          divisor_override=divisor_override,
      )
  
  
  @register_acc_op_mapping(op_and_target=("call_function", nn.functional.avg_pool3d))
  @register_acc_op
  def avg_pool3d(
      *,
      input,
      kernel_size,
      stride,
      padding,
      ceil_mode,
      count_include_pad,
      divisor_override,
  ):
      return nn.functional.avg_pool3d(
          input=input,
          kernel_size=kernel_size,
          stride=stride,
          padding=padding,
          ceil_mode=ceil_mode,
          count_include_pad=count_include_pad,
          divisor_override=divisor_override,
      )
  
  
  @register_acc_op_properties(AccOpProperty.pointwise, AccOpProperty.unary)
  @register_acc_op_mapping(op_and_target=("call_function", torch.sign))
  @register_acc_op
  def sign(*, input):
      return torch.sign(input)
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_method", "type"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("dtype", "dtype", this_arg_is_optional),
      ],
  )
  def custom_type_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
      input_obj = node.kwargs["input"]
      dtype_obj = node.kwargs.get("dtype")
      with node.graph.inserting_before(node):
          if dtype_obj is None:
              dtype_node = node.graph.call_function(dtype, kwargs={"input": input_obj})
              dtype_node.meta["type"] = torch.dtype
              return dtype_node
          else:
              new_kwargs = {
                  "input": input_obj,
                  "acc_out_ty": acc_utils.build_raw_tensor_meta(dtype=dtype_obj),
              }
              new_node = node.graph.call_function(to_dtype, kwargs=new_kwargs)
              new_node.meta = node.meta
              return new_node
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_method", "type_as"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("tensor", "tensor"),
      ],
  )
  def custom_type_as_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
      input_obj = node.kwargs["input"]
      other_obj = node.kwargs["tensor"]
      with node.graph.inserting_before(node):
          dtype_node = node.graph.call_function(dtype, kwargs={"input": other_obj})
          dtype_node.meta["type"] = torch.dtype
          device_node = node.graph.call_function(device, kwargs={"input": other_obj})
          device_node.meta["type"] = torch.device
  
          new_kwargs = {
              "input": input_obj,
              "acc_out_ty": acc_utils.build_raw_tensor_meta(dtype=dtype_node),
              "device": device_node,
          }
          new_node = node.graph.call_function(to_dtype, kwargs=new_kwargs)
          new_node.meta = node.meta
          return new_node
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op
  def dtype(*, input):
      return input.dtype
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op
  def size(*, input):
      return input.size()
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op
  def device(*, input):
      return input.device
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op_mapping(op_and_target=("call_function", torch.numel))
  @register_acc_op
  def numel(*, input):
      return torch.numel(input)
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", getattr),
      arg_replacement_tuples=[],
  )
  def custom_getattr_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
      """
      Custom function for mapping a call_function getattr to other ops.
  
      Supports:
      * getattr on a torch.Tensor with "shape", "device", or "dtype" attributes
      * getattr for accessing named tuples
      """
      # Have to use args here since getattr forces positional args.
      input_obj = node.args[0]
      attr_name = node.args[1]
      assert isinstance(input_obj, torch.fx.Node)
      input_obj_type = input_obj.meta["type"]
  
      # Handle named tuple access. NamedTupleMeta and the namedtuple factory function
      # create a subclass of tuple with an extra _fields attribute.
      if issubclass(input_obj_type, tuple) and hasattr(input_obj_type, "_fields"):
          idx = None
          for i, name in enumerate(input_obj_type._fields):
              if name == attr_name:
                  idx = i
                  break
          assert (
              idx is not None
          ), f"Named tuple type {input_obj_type} does not have field {name}"
  
          with node.graph.inserting_before(node):
              getitem_node = node.graph.call_function(
                  getitem, kwargs={"input": input_obj, "idx": idx}
              )
              getitem_node.meta = node.meta.copy()
              return getitem_node
  
      assert input_obj_type in [
          torch.Tensor,
          torch.nn.parameter.Parameter,
      ], f"Expected torch.Tensor type for {input_obj_type}"
      assert (
          attr_name == "shape" or attr_name == "device" or attr_name == "dtype"
      ), f"Only supporting shape, device and dtype getattr for now, not {attr_name}"
      if attr_name == "shape":
          func = size
      elif attr_name == "device":
          func = device
      elif attr_name == "dtype":
          func = dtype
      with node.graph.inserting_before(node):
          size_node = node.graph.call_function(func, kwargs={"input": input_obj})
          size_node.meta = node.meta.copy()
          return size_node
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_method", "size"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("dim", "dim", this_arg_is_optional),
      ],
  )
  def tensor_size_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
      """
      Mapping from Tensor.size() to acc_ops.size. We map size() to acc_ops.size directly
      and map size(dim) to acc_ops.size + acc_ops.getitem.
      """
  
      with node.graph.inserting_before(node):
          size_node = node.graph.call_function(
              size, kwargs={"input": node.kwargs["input"]}
          )
  
          if "dim" not in node.kwargs:
              size_node.meta = node.meta.copy()
              return size_node
  
          size_node.meta["type"] = torch.Size
          getitem_node = node.graph.call_function(
              getitem, kwargs={"input": size_node, "idx": node.kwargs["dim"]}
          )
          getitem_node.meta = node.meta.copy()
          return getitem_node
  
  
  @register_acc_op_properties(AccOpProperty.pointwise)
  @register_acc_op_mapping(op_and_target=("call_function", operator.add))
  @register_acc_op_mapping(op_and_target=("call_method", "add"))
  @register_acc_op
  def add(*, input, other):
      if not (isinstance(input, torch.Tensor) or isinstance(other, torch.Tensor)):
          return operator.add(input, other)
      else:
          return input + other
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op_mapping(op_and_target=("call_method", "unsqueeze"))
  @register_acc_op_mapping(op_and_target=("call_function", torch.unsqueeze))
  @register_acc_op
  def unsqueeze(*, input, dim: int):
      return torch.unsqueeze(input=input, dim=dim)
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op_mapping(op_and_target=("call_method", "tile"))
  @register_acc_op_mapping(op_and_target=("call_function", torch.tile))
  @register_acc_op
  def tile(*, input, dims):
      return torch.tile(input=input, dims=dims)
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_method", "repeat"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("*", "sizes"),
      ],
      skip_normalization_if_none=True,
  )
  def repeat_mapper(node: torch.fx.Node, _: nn.Module) -> Optional[torch.fx.Node]:
      """
      Map repeat to tile.
      """
      with node.graph.inserting_before(node):
          inputs = node.kwargs["input"]
          dims = node.kwargs["sizes"]
          # Skip repeat mapping when the list of dims is not all ints (ie. contains
          # some calculated value). torch.tile cannot support cases where dims
          # are Proxy nodes
          if (
              isinstance(dims, (list, tuple))
              and len(dims) > 0
              and not all(isinstance(x, int) for x in dims)
          ):
              logger.info(
                  "Not mapping repeat to an acc op. We can't handle variable dims."
              )
              return
          new_node = node.graph.create_node(
              "call_function",
              tile,
              kwargs={"input": inputs, "dims": dims},
              name=f"{node.name}_repeat_map",
          )
          new_node.meta = node.meta.copy()
          return new_node
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_method", "repeat_interleave"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("repeats", "repeats"),
          ("dim", "dim", this_arg_is_optional),
          ("output_size", "output_size", this_arg_is_optional),
      ],
      skip_normalization_if_none=True,
  )
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", torch.repeat_interleave),
      arg_replacement_tuples=[
          ("input", "input"),
          ("repeats", "repeats"),
          ("dim", "dim", this_arg_is_optional),
          ("output_size", "output_size", this_arg_is_optional),
      ],
      skip_normalization_if_none=True,
  )
  def repeat_interleave_mapper(node: torch.fx.Node, _: nn.Module):
      input_node = node.kwargs["input"]
      repeats = cast(int, node.kwargs["repeats"])
      dim = node.kwargs["dim"]
      if not (type(repeats) is int):
          logger.info(
              "Not mapping repeat_interleave to an acc op. We currently only support `repeat_interleave` with int repeats"
          )
          return
      rank = node.meta["tensor_rank"]
      if dim is None:
          repeat_dim = rank - 1
      else:
          assert type(dim) is int, "dim should be an int"
          repeat_dim = dim
      tile_dims = [1] * (rank + 1)
      tile_dims[repeat_dim + 1] = repeats
  
      with node.graph.inserting_before(node):
          unsqueeze_node = node.graph.create_node(
              "call_function",
              unsqueeze,
              kwargs={"input": input_node, "dim": repeat_dim + 1},
              name=f"{node.name}_unsqueeze",
          )
          tile_node = node.graph.create_node(
              "call_function",
              tile,
              kwargs={"input": unsqueeze_node, "dims": tuple(tile_dims)},
              name=f"{node.name}_repeat_interleave_map_tile",
          )
          new_shape = []
          if dim is not None:
              if dim < 0:
                  repeat_dim = dim + rank
              else:
                  repeat_dim = dim
              size_node = node.graph.create_node(
                  "call_function",
                  size,
                  kwargs={"input": input_node},
                  name=f"{node.name}_size",
              )
              size_node.meta["type"] = torch.Size
              for i in range(rank):
                  shape_i = node.graph.create_node(
                      "call_function",
                      getitem,
                      kwargs={"input": size_node, "idx": i},
                      name=f"{node.name}_size_{i}",
                  )
                  if i == repeat_dim:
                      new_shape.append(-1)
                  else:
                      new_shape.append(shape_i)
          else:
              new_shape.append(-1)
  
          reshaped_node = node.graph.create_node(
              "call_function",
              reshape,
              kwargs={
                  "input": tile_node,
                  "acc_out_ty": acc_utils.build_raw_tensor_meta(shape=new_shape),
              },
              name=f"{node.name}_reshape",
          )
          reshaped_node.meta = node.meta.copy()
          return reshaped_node
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", torch.stack),
      arg_replacement_tuples=[
          ("tensors", "tensors"),
          ("dim", "dim"),
      ],
  )
  def stack_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
      """
      Map torch.stack to unsqueeze + cat.
      """
      with node.graph.inserting_before(node):
          inputs = node.kwargs["tensors"]
          unsqueeze_nodes = []
          assert isinstance(inputs, Sequence)
          for i, t in enumerate(inputs):
              new_node = node.graph.create_node(
                  "call_function",
                  unsqueeze,
                  kwargs={"input": t, "dim": node.kwargs["dim"]},
                  name=f"{node.name}_unsqueeze_{i}",
              )
              new_node.meta["type"] = torch.Tensor
              unsqueeze_nodes.append(new_node)
          cat_node = node.graph.create_node(
              "call_function",
              cat,
              kwargs={"tensors": unsqueeze_nodes, "dim": node.kwargs["dim"]},
          )
          cat_node.meta = node.meta.copy()
          return cat_node
  
  
  @register_acc_op_properties(AccOpProperty.pointwise)
  @register_acc_op_mapping(op_and_target=("call_function", torch.clamp))
  @register_acc_op_mapping(op_and_target=("call_function", torch.clip))
  @register_acc_op_mapping(op_and_target=("call_method", "clamp"))
  @register_acc_op_mapping(op_and_target=("call_method", "clip"))
  @register_acc_op
  def clamp(*, input, min=None, max=None):
      return torch.clamp(input=input, min=min, max=max)
  
  
  @register_acc_op_mapping(op_and_target=("call_function", torch.concat))
  @register_acc_op_mapping(op_and_target=("call_function", torch.cat))
  @register_acc_op
  def cat(*, tensors, dim):
      return torch.cat(tensors=tensors, dim=dim)
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", torch.transpose),
      arg_replacement_tuples=[
          ("input", "input"),
          ("dim0", "dim0"),
          ("dim1", "dim1"),
      ],
  )
  @register_custom_acc_mapper_fn(
      op_and_target=("call_method", "transpose"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("dim0", "dim0"),
          ("dim1", "dim1"),
      ],
  )
  def transpose_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
      # Get the dim-permutation/shuffle
      ranks = node.meta["tensor_rank"]
      shuffle = list(range(ranks))
      dim0 = cast(int, node.kwargs["dim0"])
      dim1 = cast(int, node.kwargs["dim1"])
      shuffle[dim0] = dim1
      shuffle[dim1] = dim0
  
      # Create the new acc_ops.permute node. Update all uses of the transpose
      # node and then delete the transpose node.
      with node.graph.inserting_after(node):
          permute_node = node.graph.call_function(
              the_function=permute,
              kwargs={
                  "input": node.kwargs.get("input"),
                  "permutation": shuffle,
              },
          )
          permute_node.meta = node.meta.copy()
          node.replace_all_uses_with(permute_node)
  
      permute_node.graph.erase_node(node)
      return permute_node
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op_mapping(op_and_target=("call_method", "contiguous"))
  @register_acc_op
  def contiguous(*, input):
      return input.contiguous()
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op_mapping(
      op_and_target=("call_method", "softmax"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("dim", "dim"),
          ("dtype", "dtype", this_arg_is_optional),
      ],
  )
  @register_acc_op_mapping(
      op_and_target=("call_function", torch.nn.functional.softmax),
      arg_replacement_tuples=[
          ("input", "input"),
          ("dim", "dim"),
          ("dtype", "dtype", this_arg_is_optional),
      ],
  )
  @register_acc_op
  def softmax(*, input, dim, dtype=None):
      """
      _stacklevel are ignored here.
      """
      return torch.nn.functional.softmax(input=input, dim=dim, dtype=dtype)
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", torch.addmm),
      arg_replacement_tuples=[
          ("input", "input"),
          ("mat1", "mat1"),
          ("mat2", "mat2"),
          ("beta", "beta"),
          ("alpha", "alpha"),
      ],
  )
  def addmm_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
      """
      Mapping from torch.addmm to acc_ops.mm -> acc_ops.add, if alpha or beta is not 1
      then we also insert acc_ops.mul to the right place.
      """
      with node.graph.inserting_before(node):
          mm_kwargs = {"input": node.kwargs["mat1"], "other": node.kwargs["mat2"]}
          mm_node = node.graph.create_node(
              "call_function", matmul, kwargs=mm_kwargs, name=f"{node.name}_mm"
          )
          mm_node.meta = node.meta.copy()
  
          if node.kwargs["alpha"] != 1:
              mul_kwargs = {"input": mm_node, "other": node.kwargs["alpha"]}
              mm_node = node.graph.create_node(
                  "call_function", mul, kwargs=mul_kwargs, name=f"{mm_node.name}_mul"
              )
          mm_node.meta = node.meta.copy()
  
          input_node = node.kwargs["input"]
          if node.kwargs["beta"] != 1:
              mul_kwargs = {"input": input_node, "other": node.kwargs["beta"]}
              new_input_node = node.graph.create_node(
                  "call_function", mul, kwargs=mul_kwargs, name=f"{node.name}_input_mul"
              )
              assert isinstance(input_node, torch.fx.Node)
              new_input_node.meta = input_node.meta.copy()
              input_node = new_input_node
  
          add_kwargs = {"input": mm_node, "other": input_node}
          add_node = node.graph.create_node(
              "call_function", add, kwargs=add_kwargs, name=f"{node.name}_add"
          )
          add_node.meta = node.meta.copy()
          return add_node
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", torch.t),
      arg_replacement_tuples=[
          ("input", "input"),
      ],
  )
  @register_custom_acc_mapper_fn(
      op_and_target=("call_method", "t"),
      arg_replacement_tuples=[
          ("input", "input"),
      ],
  )
  def t_mapper(node: torch.fx.Node, _: nn.Module):
      ranks = node.meta["tensor_rank"]
      shuffle = [1, 0] if (ranks > 1) else [0]
  
      with node.graph.inserting_before(node):
          new_node = node.graph.create_node(
              "call_function",
              permute,
              kwargs={"input": node.kwargs["input"], "permutation": shuffle},
          )
          new_node.meta = node.meta.copy()
          return new_node
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op_mapping(
      op_and_target=("call_method", "permute"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("*", "permutation"),
      ],
  )
  @register_acc_op_mapping(
      op_and_target=("call_function", torch.permute),
      arg_replacement_tuples=[
          ("input", "input"),
          ("dims", "permutation"),
      ],
  )
  @register_acc_op
  def permute(*, input, permutation):
      return input.permute(*permutation)
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", torch.square),
      arg_replacement_tuples=[
          ("input", "input"),
      ],
  )
  def square_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
      input_node = node.kwargs["input"]
      with node.graph.inserting_before(node):
          new_node = node.graph.call_function(
              mul, kwargs={"input": input_node, "other": input_node}
          )
          new_node.meta = node.meta.copy()
          return new_node
  
  
  @register_acc_op_mapping(
      op_and_target=("call_method", "mm"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("mat2", "other"),
      ],
  )
  @register_acc_op_mapping(
      op_and_target=("call_method", "matmul"),
      arg_replacement_tuples=[
          ("input", "input"),
          ("mat2", "other"),
      ],
  )
  @register_acc_op_mapping(
      op_and_target=("call_function", operator.matmul),
      arg_replacement_tuples=[
          ("input", "input"),
          ("mat2", "other"),
      ],
  )
  @register_acc_op_mapping(
      op_and_target=("call_function", torch.bmm),
      arg_replacement_tuples=[
          ("input", "input"),
          ("mat2", "other"),
      ],
  )
  @register_acc_op_mapping(
      op_and_target=("call_function", torch.mm),
      arg_replacement_tuples=[
          ("input", "input"),
          ("mat2", "other"),
      ],
  )
  @register_acc_op_mapping(op_and_target=("call_function", torch.matmul))
  @register_acc_op
  def matmul(*, input, other):
      return torch.matmul(input=input, other=other)
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", nn.functional.dropout),
      arg_replacement_tuples=[("input", "input")],
  )
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", nn.functional.dropout1d),
      arg_replacement_tuples=[("input", "input")],
  )
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", nn.functional.dropout2d),
      arg_replacement_tuples=[("input", "input")],
  )
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", nn.functional.dropout3d),
      arg_replacement_tuples=[("input", "input")],
  )
  @register_custom_acc_mapper_fn(
      op_and_target=("call_method", "detach"), arg_replacement_tuples=[("input", "input")]
  )
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", torch.detach),
      arg_replacement_tuples=[("input", "input")],
  )
  def dropout_mapper(node: torch.fx.Node, mod: nn.Module):
      """
      Remove dropout node and directly map its input to output.
      """
      return node.kwargs["input"]
  
  
  try:
      from torchvision.ops import stochastic_depth
  
      assert callable(stochastic_depth)
  except Exception as e:
      warnings.warn(
          f"Unable to import torchvision related libraries.: {e}. Please install torchvision lib in order to lower stochastic_depth"
      )
  else:
  
      @register_custom_acc_mapper_fn(
          op_and_target=("call_function", stochastic_depth),
          arg_replacement_tuples=[("input", "input")],
      )
      def stochastic_depth_mapper(node: torch.fx.Node, mod: nn.Module):
          """
          Remove dropout node and directly map its input to output.
          """
          return node.kwargs["input"]
  
  
  @register_acc_op_properties(AccOpProperty.pointwise, AccOpProperty.unary)
  @register_acc_op_mapping(
      op_and_target=("call_function", nn.functional.hardtanh),
  )
  @register_acc_op
  def hardtanh(*, input, min_val=-1.0, max_val=1.0):
      return nn.functional.hardtanh(input=input, min_val=min_val, max_val=max_val)
  
  
  @register_acc_op_properties(AccOpProperty.pointwise, AccOpProperty.unary)
  @register_acc_op_mapping(op_and_target=("call_function", nn.functional.hardsigmoid))
  @register_acc_op
  def hardsigmoid(*, input):
      return nn.functional.hardsigmoid(input)
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", nn.functional.silu),
      arg_replacement_tuples=[
          ("input", "input"),
      ],
  )
  def silu(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
      input_node = node.kwargs["input"]
      with node.graph.inserting_before(node):
          sigmoid_node = node.graph.call_function(sigmoid, kwargs={"input": input_node})
          sigmoid_node.meta = node.meta.copy()
          new_node = node.graph.call_function(
              mul, kwargs={"input": sigmoid_node, "other": input_node}
          )
          new_node.meta = node.meta.copy()
          return new_node
  
  
  @register_custom_acc_mapper_fn(
      op_and_target=("call_function", nn.functional.hardswish),
      arg_replacement_tuples=[
          ("input", "input"),
      ],
  )
  def hardswish_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
      input_node = node.kwargs["input"]
      with node.graph.inserting_before(node):
          new_sigmoid_node = node.graph.call_function(
              hardsigmoid, kwargs={"input": input_node}
          )
          new_sigmoid_node.meta = node.meta.copy()
          new_node = node.graph.call_function(
              mul, kwargs={"input": new_sigmoid_node, "other": input_node}
          )
          new_node.meta = node.meta.copy()
          return new_node
  
  
  @register_acc_op_properties(AccOpProperty.quantized)
  @register_acc_op_mapping(
      op_and_target=("call_function", torch.ops.quantized.add),
      arg_replacement_tuples=[
          ("qa", "input"),
          ("qb", "other"),
          ("scale", "scale"),
          ("zero_point", "zero_point"),
      ],
      kwargs_to_move_to_acc_out_ty=[
          ("scale", "scale", move_to_qparams),
          ("zero_point", "zero_point", move_to_qparams),
      ],
  )
  @register_acc_op
  def quantized_add(*, input, other, acc_out_ty=None):
      assert acc_out_ty is not None
      qparams = acc_out_ty.qparams
      return torch.ops.quantized.add(
          input,
          other,
          qparams["scale"],
          qparams["zero_point"],
      )
  
  
  @register_acc_op_properties(AccOpProperty.quantized)
  @register_acc_op_mapping(
      op_and_target=("call_function", torch.ops.quantized.mul),
      arg_replacement_tuples=[
          ("qa", "input"),
          ("qb", "other"),
          ("scale", "scale"),
          ("zero_point", "zero_point"),
      ],
      kwargs_to_move_to_acc_out_ty=[
          ("scale", "scale", move_to_qparams),
          ("zero_point", "zero_point", move_to_qparams),
      ],
  )
  @register_acc_op
  def quantized_mul(*, input, other, acc_out_ty=None):
      assert acc_out_ty is not None
      qparams = acc_out_ty.qparams
      return torch.ops.quantized.mul(
          input,
          other,
          qparams["scale"],
          qparams["zero_point"],
      )
  
  
  @register_acc_op_properties(AccOpProperty.pointwise, AccOpProperty.unary)
  @register_acc_op_properties(AccOpProperty.quantized)
  @register_acc_op_mapping(
      op_and_target=("call_function", torch.quantize_per_tensor),
      arg_replacement_tuples=[
          ("input", "input"),
          ("scale", "scale"),
          ("zero_point", "zero_point"),
          ("dtype", "dtype"),
      ],
      kwargs_to_move_to_acc_out_ty=[
          ("scale", "scale", move_to_qparams),
          ("zero_point", "zero_point", move_to_qparams),
          ("dtype", "dtype", dont_move_to_qparams),
      ],
  )
  @register_acc_op
  def quantize_per_tensor(*, input, acc_out_ty=None):
      assert acc_out_ty is not None
      qparams = acc_out_ty.qparams
      dtype = acc_out_ty.dtype
      return torch.quantize_per_tensor(
          input, qparams["scale"], qparams["zero_point"], dtype
      )
  
  
  @register_acc_op_properties(AccOpProperty.unary)
  @register_acc_op_mapping(
      op_and_target=("call_function", torch.quantize_per_channel),
      arg_replacement_tuples=[
          ("input", "input"),
          ("scales", "scales"),
          ("zero_points", "zero_points"),
          ("axis", "axis"),
          ("dtype", "dtype"),
      ],
      kwargs_to_move_to_acc_out_ty=[
          ("scales", "scale", move_to_qparams),
          ("zero_points", "zero_point", move_to_qparams),
          ("axis", "axis", move_to_qparams),
          ("dtype", "dtype", dont_move_to_qparams),
      ],
  )
  @register_acc_op
  def quantize_per_channel(*, input, acc_out_ty=None):
      assert acc_out_ty is not None
      qparams = acc_out_ty.qparams
      dtype = acc_out_ty.dtype
      return torch.quantize_per_channel(
          input,
          torch.tensor(qparams["scale"]),
          torch.tensor(qparams["zero_point"]),
          qparams["axis"],
          dtype,
      )  # type: ignore[call-overload]
  
  
  @register_acc_op_properties(AccOpProperty.pointwise, AccOpProperty.unary)
  @register_acc_op_mapping(op_and_target=("call_method", "dequantize"))
  @register_acc_op_mapping(op_and_target=("call_function", torch.dequantize))
  @register_acc_op
  def dequantize(*, input):
      return torch.dequantize(input)
  
  
  @register_acc_op_properties(
      AccOpProperty.pointwise, AccOpProperty.unary, AccOpProperty.quantized
  )
  @register_acc_op
  def rescale_quantize_per_tensor(*, input, acc_out_ty=None):
  