content: from typing import Any, Callable, Dict, Set, Union
  
  import torch
  from torch._decomp import core_aten_decompositions
  from torch._decomp import get_decompositions as get_torch_decompositions
  from torch._ops import OpOverload, OpOverloadPacket
  
  aten = torch.ops.aten
  
  _core_aten_decompositions: Dict[OpOverload, Callable[[Any], Any]] = (
      core_aten_decompositions()
  )
  torch_enabled_decompositions: Set[Union[OpOverload, OpOverloadPacket]] = {
      aten._adaptive_avg_pool2d_backward,
      aten.addcdiv,
      aten.addcdiv_,
      aten.addcmul,
      aten.addcmul_,
      aten.addr,
      aten.aminmax,
      aten.arange.default,
      aten.arange.start,
      aten.avg_pool2d_backward,
      aten.binary_cross_entropy,
      aten.binary_cross_entropy_backward,
      aten.binary_cross_entropy_with_logits,
      aten.celu,
      aten.col2im,
      aten.count_nonzero,
      aten.cudnn_batch_norm,
      aten.cudnn_batch_norm_backward,
      aten.deg2rad,
      aten.detach,
      aten.diag_embed,
      aten.diagonal_backward,
      aten.elu_backward,
      aten.embedding_dense_backward,
      aten.empty_like,
      aten._euclidean_dist.default,
      aten.expand_as,
      aten.eye,
      aten.fill,
      aten.frac,
      aten._fused_moving_avg_obs_fq_helper,
      aten.gelu_backward,
      aten.glu_backward,
      aten.hardshrink,
      aten.hardshrink_backward,
      aten.hardsigmoid_backward,
      aten.hardswish,
      aten.hardswish_,
      aten.hardswish_backward,
      aten.hardtanh,
      aten.hardtanh_,
      aten.hardtanh_backward,
      aten.heaviside,
      aten.huber_loss,
      aten.huber_loss_backward,
      aten.im2col,
      aten.index_add,
      aten.index_add_,
      aten.index_copy,
      aten.index_copy_,
      aten.index_fill,
      aten.index_fill_,
      aten.isneginf,
      aten.isposinf,
      aten.l1_loss,
      aten.leaky_relu_,
      aten.leaky_relu_backward,
      aten.lerp,
      aten.linspace,
      aten.logaddexp,
      aten.logaddexp2,
      aten.logit,
      aten.logit_backward,
      aten.log_sigmoid_backward,
      aten.log_sigmoid_forward,
      aten._log_softmax_backward_data,
      aten.logspace,
      aten.logsumexp.default,
      aten.masked_fill,
      aten.masked_fill_,
      aten.max_pool2d_with_indices_backward,
      aten.mish,
      aten.mse_loss,
      aten.mse_loss_backward,
      aten.mvlgamma,
      aten.nansum,
      aten.nan_to_num,
      aten.narrow,
      # TODO: Disable the below operators once freezing is done
      aten.native_batch_norm_backward,
      aten._native_batch_norm_legit_functional,
      aten.native_dropout_backward,
      aten.native_group_norm_backward,
      aten.native_layer_norm_backward,
      aten.new_empty,
      aten.new_full,
      aten.new_ones,
      aten.new_zeros,
      aten.nll_loss_backward,
      aten.nll_loss_forward,
      aten.norm,
      aten.ones,
      aten.ones_like,
      aten._prelu_kernel_backward,
      aten._reshape_alias,
      aten.rad2deg,
      aten.renorm,
      aten.renorm_,
      aten.rot90,
      aten.rsub,
      aten.select_backward,
      aten.select_scatter,
      aten.sgn,
      aten.sigmoid_backward,
      aten.silu_backward,
      aten.sinc,
      aten.slice_backward,
      aten.smooth_l1_loss,
      aten.smooth_l1_loss_backward,
      aten.soft_margin_loss,
      aten.soft_margin_loss_backward,
      aten._softmax.out,
      aten._softmax_backward_data,
      aten.softplus_backward,
      aten.softshrink,
      aten.softshrink_backward,
      aten.special_entr,
      aten.special_log_ndtr,
      aten.special_xlog1py,
      aten.stack,
      aten.std,
      aten.t,
      aten.tanh_backward,
      aten.threshold,
      aten.threshold_backward,
      aten.trace,
      aten.transpose.int,
      aten.tril.default,
      aten.triu.default,
      aten.unbind,
      aten.unfold,
      aten.unfold_backward,
      aten.unfold_copy,
      aten._unsafe_index,
      aten.upsample_nearest2d_backward,
      aten.var,
      aten.var_mean,
      aten.xlogy,
      aten.zero,
      aten.zero_,
      aten.zeros,
      aten.zeros_like,
      # Non-default convenience decompositions
      aten.clamp_min,
      aten.clamp_max,
      aten.linalg_vector_norm,
      aten.repeat,
  }
  torch_disabled_decompositions: Set[Union[OpOverload, OpOverloadPacket]] = {
      aten._softmax.default,
      aten.upsample_nearest1d.vec,
      aten.upsample_nearest2d.vec,
      aten.upsample_nearest3d.vec,
      aten.upsample_linear1d.vec,
      aten.upsample_bilinear2d.vec,
      aten.upsample_trilinear3d.vec,
      aten.upsample_bicubic2d.vec,
      aten.linear.default,
      aten.matmul.default,
  }
  
  
  ENABLED_TORCH_DECOMPOSITIONS: Dict[OpOverload, Callable[[Any], Any]] = (
      get_torch_decompositions(torch_enabled_decompositions)
  )
  TORCH_TRT_DECOMPOSITIONS: Dict[OpOverload, Callable[[Any], Any]] = {}
  
  
  def check_decomp_set_invariants() -> None:
      """Validates no overlap between enabled and disabled decomposition sets"""
      overlap = torch_enabled_decompositions.intersection(torch_disabled_decompositions)
  
      if overlap:
          raise AssertionError(
              f"Detected {overlap} registered in both torch_enabled_decompositions "
              "and torch_disabled_decompositions. Ensure all operator(s) are in "
              "at most one of the two sets."
          )
  
  
  check_decomp_set_invariants()
  