# LANGUAGE AllowAmbiguousTypes , ConstraintKinds , DataKinds , ExtendedDefaultRules , FlexibleContexts , FlexibleInstances , GADTs         #
# LANGUAGE KindSignatures , MultiParamTypeClasses , OverloadedStrings , PolyKinds , QuasiQuotes , RankNTypes , ScopedTypeVariables         #
# LANGUAGE TemplateHaskell , TypeApplications , TypeFamilies , TypeFamilyDependencies , TypeInType , TypeOperators , UndecidableInstances #
# OPTIONS_GHC -fplugin GHC.TypeLits . KnownNat . Solver #

module Torch.Visualization where
import           Control.Monad
import qualified Data.List            as L
import qualified Data.Map.Strict      as M
import           Data.Text            (Text)
import qualified Data.Text            as T
import qualified Data.Text.IO         as T
import           Data.Vector.Storable (Vector)
import qualified Data.Vector.Storable as V
import           Foreign.ForeignPtr
import           Foreign.Ptr
import           GHC.Int
import qualified Shelly               as S
import qualified Torch.C.Tensor       as C
import qualified Torch.C.Types        as C
import qualified Torch.C.Variable     as C
import           Torch.Tensor
import           Torch.Types
import           Foreign.ForeignPtr.Unsafe
import           Torch.Misc

data TraceValue = TraceValueTensor { tvTensorName       :: T.Text
                                   , tvTensorSizes      :: Vector Int64
                                   , tvtensorScalarType :: C.ScalarType }
                deriving (Show, Eq, Ord)

data TraceNode = TraceNode { tnKind    :: T.Text
                           , tnInputs  :: [TraceValue]
                           , tnOutputs :: [TraceValue] }
                deriving (Show, Eq, Ord)

data TraceGraph = TraceGraph { traceGraphInputs  :: [TraceValue]
                             , traceGraphNodes   :: [TraceNode]
                             , traceGraphOutputs :: [TraceValue] }
                deriving (Show, Eq, Ord)

type TracePtr = ForeignPtr C.CTracingState

withTracing :: (TensorConstraints ty ki sz) => [AnyTensor] -> IO (Tensor ty ki sz) -> IO (Tensor ty ki sz, TracePtr)
withTracing ts f = do
  TODO This has a constant , 1
  cfn <- C.mkTraceable (\_ -> castPtr . unsafeForeignPtrToPtr . tensorPtr <$> f) -- TODO Tensor to Variable cast
  (state, vars) <- withForeignPtrList (map (\(AnyTensor (Tensor p _)) -> p) ts)
                                      (\ps -> C.trace (V.map castPtr ps) 1 cfn)
  -- FIXME What about this?
  state ' < - newForeignPtr C.deleteTracingState state
  state' <- newForeignPtr_ state
  ptr <- newForeignPtr C.deleteTensor  $ castPtr $ vars V.! 0
  pure (Tensor ptr Nothing, state') -- TODO Nothing might instead be vars

-- | This prints the trace in the pytorch/onnx graph format.
printTrace :: TracePtr -> IO ()
printTrace s = withForeignPtr s C.print_tracing_state_graph

-- | This prints the trace in the onnx graph format.
printTraceONNX :: TracePtr -> [AnyTensor] -> Bool -> Int -> IO ()
printTraceONNX s ts googleFormat opset =
  withForeignPtrList (map (\(AnyTensor (Tensor p _)) -> p) ts)
  (\ps -> withForeignPtr s (\s -> C.print_tracing_state_graph_onnx s (V.map castPtr ps) (boolc googleFormat) (fromIntegral opset)))

tnIsConstant :: TraceNode -> Bool
tnIsConstant tn | tnKind tn == "prim::Constant" = True
                | otherwise = False

showsz :: (V.Storable a, Show a) => Vector a -> Text
showsz x =
  case V.length x of
    0 -> "Scalar"
    _ -> T.pack $ show x

data TraceType = TraceTypeSimple Bool C.TypeKind
               | TraceTypeList   Bool [TraceType]
               | TraceTypeTuple  Bool [TraceType]
               | TraceTypeTensor Bool (Maybe (Vector Int64)) (Maybe C.ScalarType) (Maybe C.Backend)
               deriving (Show, Eq)

convertCType :: Ptr C.CType -> IO TraceType
convertCType t = do
  k <- C.type_kind t
  grad <- C.type_requires_grad t
  case k of
    C.TypeKindTensor -> do
      sizes <- C.type_sizes t
      scalar <- C.type_scalar_type t
      device <- C.type_device_type t
      pure $ TraceTypeTensor grad (Just sizes) (Just scalar) (Just device)
  TODO cleanup
      -- C.TypeKindCompleteTensor -> do
      -- pure $ TraceTypeTensor grad Nothing Nothing Nothing
    C.TypeKindList -> do
      cs <- C.type_contained t
      cs' <- mapM convertCType $ V.toList cs
      pure $ TraceTypeList grad cs'
    C.TypeKindTuple -> do
      cs <- C.type_contained t
      cs' <- mapM convertCType $ V.toList cs
      pure $ TraceTypeTuple grad cs'
    C.TypeKindInt -> do
      pure $ TraceTypeSimple grad k
    C.TypeKindFloat -> do
      pure $ TraceTypeSimple grad k
    C.TypeKindString -> do
      pure $ TraceTypeSimple grad k
    C.TypeKindNone -> do
      pure $ TraceTypeSimple grad k
    C.TypeKindGenerator -> do
      pure $ TraceTypeSimple grad k
    C.TypeKindBool -> do
      pure $ TraceTypeSimple grad k
    _ -> do
      print =<< C.type_string t
      error "Unsupported variable type, see line above"

parseTrace :: ForeignPtr C.CTracingState -> IO TraceGraph
parseTrace trace =
  parseGraph =<< withForeignPtr trace C.tracing_state_graph
  where parseGraph (inputs, nodes, outputs, _, _) = do
          is <- mapM parseValue $ V.toList inputs
          ns <- mapM parseNode $ V.toList nodes
          os <- mapM parseValue $ V.toList outputs
          pure $ TraceGraph is ns os
        parseValue x = do
          b <- cbool <$> C.check_value_tensor x
          if b then do
            sz <- C.value_sizes x
            ty <- C.value_scalar_type x
            na <- C.value_name x
            pure $ TraceValueTensor (T.pack na) sz ty
            else do
            -- TODO Figure out more stuff here
            na <- C.value_name x
            pure $ TraceValueTensor (T.pack na) V.empty C.ScalarTypeUndefined
            -- DEBUGGING TODO
            -- error "Other values are not yet supported"
        parseNode x = do
          ki <- C.node_kind x
          is <- mapM parseValue . V.toList =<< C.node_inputs x
          os <- mapM parseValue . V.toList =<< C.node_outputs x
          pure $ TraceNode (T.pack ki) is os

countTraceParameters :: TraceGraph -> Int64
countTraceParameters trace = countGraph trace
  where countGraph tr = Prelude.sum $ map countNode $ traceGraphNodes tr
        countNode n = if tnIsConstant n then
                        Prelude.sum (map countValue (tnOutputs n))
                        else 0
        countValue v = case v of
          TraceValueTensor n s t -> V.product s

summarizeTrace :: TraceGraph -> IO ()
summarizeTrace trace = do
  T.putStrLn (T.justifyRight 25 ' ' "Input shape"
             <> T.justifyRight 35 ' ' (T.intercalate " " (map showsz (map tvTensorSizes (traceGraphInputs trace)))))
  T.putStrLn (T.replicate 79 "=")
  T.putStrLn (T.justifyRight 25 ' ' "Function"
             <>T.justifyRight 35 ' '"Output shape"
             <>T.justifyRight 19 ' ' "# Params")
  T.putStrLn (T.replicate 79 "-")
  mapM (\e -> do
           if tnIsConstant e then
             pure () else
             T.putStrLn
             (T.justifyRight 25 ' ' (tnKind e)
              <>T.justifyRight 35 ' '
               (T.intercalate " "
                 (map showsz (map tvTensorSizes (tnOutputs e))))
               <>T.justifyRight 19 ' '
               (T.pack
                (show
                 (Prelude.sum (map (\i ->
                                    case M.lookup (tvTensorName i) constants of
                                      Just _  -> V.product $ tvTensorSizes i
                                      Nothing -> 0)
                             (tnInputs e)))))))
    (traceGraphNodes trace)
  T.putStrLn (T.replicate 79 "=")
  T.putStrLn $ "Total number of parameters: " <> T.pack (show (countTraceParameters trace))
  T.putStrLn (T.replicate 79 "-")
  where constants = M.fromList
                    $ filter (\(_,x) -> not (elem x (traceGraphInputs trace)))
                    $ concatMap (\x ->
                           if tnIsConstant x then
                             map (\x -> (case x of
                                           TraceValueTensor{} -> tvTensorName x
                                       , x)) (tnOutputs x)
                             else [])
                    (traceGraphNodes trace)

showTraceGraph :: TracePtr -> Bool -> IO ()
showTraceGraph tptr noConstants = do
  trace <- parseTrace tptr
  displayDot (traceToDot trace noConstants)
  pure ()

data TraceGraphState = TGS { tgsStr     :: Text
                           , tgsMap     :: M.Map Text Int
                           , tgsNext    :: Int
                           , tgsOutputs :: M.Map TraceValue [Int] }
  deriving (Show)

traceToDot :: TraceGraph -> Bool -> Text
traceToDot trace noConstants =
  let state = loopGraph (TGS "" M.empty 0 M.empty) trace
  in "digraph trace {\n"
     <> "node [style=rounded]" <> "\n"
     <> tgsStr state
     <> T.unlines (map (\(v,is) -> mkOutputNode v is) (M.toList (tgsOutputs state)))
     <> "}"
  where
    withNextNode s f = f (s { tgsNext = tgsNext s + 1 }) (tgsNext s)
    mkInputNode v s i =
      s { tgsStr = tgsStr s
                 <> T.pack (show i) <> " [shape=egg,color=green,label=\"In:" <> showsz (tvTensorSizes v) <> "\"]" <> "\n"
        , tgsMap = M.insert (tvTensorName v) i (tgsMap s) }
    mkNode n s i | any (\x -> T.isPrefixOf (T.pack x) (tnKind n) || T.isSuffixOf (T.pack x) (tnKind n)) censoredNodes = s
                 | noConstants && tnIsConstant n = s
                 | otherwise =
      s { tgsStr = tgsStr s
                 <> T.pack (show i) <> " [shape=box,"
                                    <> (if tnIsConstant n then
                                          "color=orange," else
                                          "")
                                    <> "label=\""
                                    <> (if tnIsConstant n then
                                         "Const:" <> (T.intercalate " "
                                                     (map showsz (map tvTensorSizes (tnOutputs n)))) else
                                       tnKind n)
                                  <>"\"]" <> "\n"
                 <> T.unlines
                       (map (\x ->
                               case M.lookup (tvTensorName x) (tgsMap s) of
                                 Nothing ->
                                   if noConstants then
                                     "" else
                                     error $ "Can't find an input? This is a bug! " <> show x <> show n
                                 Just x -> T.pack (show x) <> " -> "  <> T.pack (show i))
                            (tnInputs n))
        , tgsMap = L.foldl' (\m v -> M.insert (tvTensorName v) i m) (tgsMap s) (tnOutputs n)
        , tgsOutputs =
            (L.foldl' (\m v ->
                       case M.lookup v m of
                         Just x  -> M.adjust (i:) v m
                         Nothing -> m)
              -- delete any seen outputs; they're outputs of a subgraph
              (L.foldl' (flip M.delete) (tgsOutputs s) (tnInputs n))
              (tnOutputs n))
        }
    mkOutputNode v is =
         "O" <> (tvTensorName v) <> " [shape=hexagon,color=blue,label=\"Out:" <> showsz (tvTensorSizes v) <> "\"]" <> "\n"
      <> T.unlines (map (\i -> T.pack (show i) <> " -> O" <> (tvTensorName v)) is)
    loopGraph state g =
      -- Invent nodes for the inputs if needed
      let state' = L.foldl' (\s v ->
                            case M.lookup (tvTensorName v) (tgsMap s) of
                              Nothing -> withNextNode s (mkInputNode v)
                              Just _  -> s)
                          state
                          (traceGraphInputs trace)
      -- Invent nodes for the outputs if needed
      in let state'' = L.foldl' (\s v -> s { tgsOutputs = M.insert v [] (tgsOutputs s) } )
                                state'
                                (traceGraphOutputs trace)
         in L.foldl' loopNode state'' (traceGraphNodes g)
    loopNode s n = withNextNode s (mkNode n)
    censoredNodes = ["aten::_convolution_nogroup"
                    ,"_forward"
                    ,"aten::thnn_"]

 TODO Tensor to Variable cast
 FIXME What about this?
 TODO Nothing might instead be vars
 | This prints the trace in the pytorch/onnx graph format.
 | This prints the trace in the onnx graph format.
 C.TypeKindCompleteTensor -> do
 pure $ TraceTypeTensor grad Nothing Nothing Nothing
 TODO Figure out more stuff here
 DEBUGGING TODO
 error "Other values are not yet supported"
 delete any seen outputs; they're outputs of a subgraph
 Invent nodes for the inputs if needed
 Invent nodes for the outputs if needed
# LANGUAGE AllowAmbiguousTypes , ConstraintKinds , DataKinds , ExtendedDefaultRules , FlexibleContexts , FlexibleInstances , GADTs         #
# LANGUAGE KindSignatures , MultiParamTypeClasses , OverloadedStrings , PolyKinds , QuasiQuotes , RankNTypes , ScopedTypeVariables         #
# LANGUAGE TemplateHaskell , TypeApplications , TypeFamilies , TypeFamilyDependencies , TypeInType , TypeOperators , UndecidableInstances #
# OPTIONS_GHC -fplugin GHC.TypeLits . KnownNat . Solver #

module Torch.Visualization where
import           Control.Monad
import qualified Data.List            as L
import qualified Data.Map.Strict      as M
import           Data.Text            (Text)
import qualified Data.Text            as T
import qualified Data.Text.IO         as T
import           Data.Vector.Storable (Vector)
import qualified Data.Vector.Storable as V
import           Foreign.ForeignPtr
import           Foreign.Ptr
import           GHC.Int
import qualified Shelly               as S
import qualified Torch.C.Tensor       as C
import qualified Torch.C.Types        as C
import qualified Torch.C.Variable     as C
import           Torch.Tensor
import           Torch.Types
import           Foreign.ForeignPtr.Unsafe
import           Torch.Misc

data TraceValue = TraceValueTensor { tvTensorName       :: T.Text
                                   , tvTensorSizes      :: Vector Int64
                                   , tvtensorScalarType :: C.ScalarType }
                deriving (Show, Eq, Ord)

data TraceNode = TraceNode { tnKind    :: T.Text
                           , tnInputs  :: [TraceValue]
                           , tnOutputs :: [TraceValue] }
                deriving (Show, Eq, Ord)

data TraceGraph = TraceGraph { traceGraphInputs  :: [TraceValue]
                             , traceGraphNodes   :: [TraceNode]
                             , traceGraphOutputs :: [TraceValue] }
                deriving (Show, Eq, Ord)

type TracePtr = ForeignPtr C.CTracingState

withTracing :: (TensorConstraints ty ki sz) => [AnyTensor] -> IO (Tensor ty ki sz) -> IO (Tensor ty ki sz, TracePtr)
withTracing ts f = do
  TODO This has a constant , 1
  (state, vars) <- withForeignPtrList (map (\(AnyTensor (Tensor p _)) -> p) ts)
                                      (\ps -> C.trace (V.map castPtr ps) 1 cfn)
  state ' < - newForeignPtr C.deleteTracingState state
  state' <- newForeignPtr_ state
  ptr <- newForeignPtr C.deleteTensor  $ castPtr $ vars V.! 0

printTrace :: TracePtr -> IO ()
printTrace s = withForeignPtr s C.print_tracing_state_graph

printTraceONNX :: TracePtr -> [AnyTensor] -> Bool -> Int -> IO ()
printTraceONNX s ts googleFormat opset =
  withForeignPtrList (map (\(AnyTensor (Tensor p _)) -> p) ts)
  (\ps -> withForeignPtr s (\s -> C.print_tracing_state_graph_onnx s (V.map castPtr ps) (boolc googleFormat) (fromIntegral opset)))

tnIsConstant :: TraceNode -> Bool
tnIsConstant tn | tnKind tn == "prim::Constant" = True
                | otherwise = False

showsz :: (V.Storable a, Show a) => Vector a -> Text
showsz x =
  case V.length x of
    0 -> "Scalar"
    _ -> T.pack $ show x

data TraceType = TraceTypeSimple Bool C.TypeKind
               | TraceTypeList   Bool [TraceType]
               | TraceTypeTuple  Bool [TraceType]
               | TraceTypeTensor Bool (Maybe (Vector Int64)) (Maybe C.ScalarType) (Maybe C.Backend)
               deriving (Show, Eq)

convertCType :: Ptr C.CType -> IO TraceType
convertCType t = do
  k <- C.type_kind t
  grad <- C.type_requires_grad t
  case k of
    C.TypeKindTensor -> do
      sizes <- C.type_sizes t
      scalar <- C.type_scalar_type t
      device <- C.type_device_type t
      pure $ TraceTypeTensor grad (Just sizes) (Just scalar) (Just device)
  TODO cleanup
    C.TypeKindList -> do
      cs <- C.type_contained t
      cs' <- mapM convertCType $ V.toList cs
      pure $ TraceTypeList grad cs'
    C.TypeKindTuple -> do
      cs <- C.type_contained t
      cs' <- mapM convertCType $ V.toList cs
      pure $ TraceTypeTuple grad cs'
    C.TypeKindInt -> do
      pure $ TraceTypeSimple grad k
    C.TypeKindFloat -> do
      pure $ TraceTypeSimple grad k
    C.TypeKindString -> do
      pure $ TraceTypeSimple grad k
    C.TypeKindNone -> do
      pure $ TraceTypeSimple grad k
    C.TypeKindGenerator -> do
      pure $ TraceTypeSimple grad k
    C.TypeKindBool -> do
      pure $ TraceTypeSimple grad k
    _ -> do
      print =<< C.type_string t
      error "Unsupported variable type, see line above"

parseTrace :: ForeignPtr C.CTracingState -> IO TraceGraph
parseTrace trace =
  parseGraph =<< withForeignPtr trace C.tracing_state_graph
  where parseGraph (inputs, nodes, outputs, _, _) = do
          is <- mapM parseValue $ V.toList inputs
          ns <- mapM parseNode $ V.toList nodes
          os <- mapM parseValue $ V.toList outputs
          pure $ TraceGraph is ns os
        parseValue x = do
          b <- cbool <$> C.check_value_tensor x
          if b then do
            sz <- C.value_sizes x
            ty <- C.value_scalar_type x
            na <- C.value_name x
            pure $ TraceValueTensor (T.pack na) sz ty
            else do
            na <- C.value_name x
            pure $ TraceValueTensor (T.pack na) V.empty C.ScalarTypeUndefined
        parseNode x = do
          ki <- C.node_kind x
          is <- mapM parseValue . V.toList =<< C.node_inputs x
          os <- mapM parseValue . V.toList =<< C.node_outputs x
          pure $ TraceNode (T.pack ki) is os

countTraceParameters :: TraceGraph -> Int64
countTraceParameters trace = countGraph trace
  where countGraph tr = Prelude.sum $ map countNode $ traceGraphNodes tr
        countNode n = if tnIsConstant n then
                        Prelude.sum (map countValue (tnOutputs n))
                        else 0
        countValue v = case v of
          TraceValueTensor n s t -> V.product s

summarizeTrace :: TraceGraph -> IO ()
summarizeTrace trace = do
  T.putStrLn (T.justifyRight 25 ' ' "Input shape"
             <> T.justifyRight 35 ' ' (T.intercalate " " (map showsz (map tvTensorSizes (traceGraphInputs trace)))))
  T.putStrLn (T.replicate 79 "=")
  T.putStrLn (T.justifyRight 25 ' ' "Function"
             <>T.justifyRight 35 ' '"Output shape"
             <>T.justifyRight 19 ' ' "# Params")
  T.putStrLn (T.replicate 79 "-")
  mapM (\e -> do
           if tnIsConstant e then
             pure () else
             T.putStrLn
             (T.justifyRight 25 ' ' (tnKind e)
              <>T.justifyRight 35 ' '
               (T.intercalate " "
                 (map showsz (map tvTensorSizes (tnOutputs e))))
               <>T.justifyRight 19 ' '
               (T.pack
                (show
                 (Prelude.sum (map (\i ->
                                    case M.lookup (tvTensorName i) constants of
                                      Just _  -> V.product $ tvTensorSizes i
                                      Nothing -> 0)
                             (tnInputs e)))))))
    (traceGraphNodes trace)
  T.putStrLn (T.replicate 79 "=")
  T.putStrLn $ "Total number of parameters: " <> T.pack (show (countTraceParameters trace))
  T.putStrLn (T.replicate 79 "-")
  where constants = M.fromList
                    $ filter (\(_,x) -> not (elem x (traceGraphInputs trace)))
                    $ concatMap (\x ->
                           if tnIsConstant x then
                             map (\x -> (case x of
                                           TraceValueTensor{} -> tvTensorName x
                                       , x)) (tnOutputs x)
                             else [])
                    (traceGraphNodes trace)

showTraceGraph :: TracePtr -> Bool -> IO ()
showTraceGraph tptr noConstants = do
  trace <- parseTrace tptr
  displayDot (traceToDot trace noConstants)
  pure ()

data TraceGraphState = TGS { tgsStr     :: Text
                           , tgsMap     :: M.Map Text Int
                           , tgsNext    :: Int
                           , tgsOutputs :: M.Map TraceValue [Int] }
  deriving (Show)

traceToDot :: TraceGraph -> Bool -> Text
traceToDot trace noConstants =
  let state = loopGraph (TGS "" M.empty 0 M.empty) trace
  in "digraph trace {\n"
     <> "node [style=rounded]" <> "\n"
     <> tgsStr state
     <> T.unlines (map (\(v,is) -> mkOutputNode v is) (M.toList (tgsOutputs state)))
     <> "}"
  where
    withNextNode s f = f (s { tgsNext = tgsNext s + 1 }) (tgsNext s)
    mkInputNode v s i =
      s { tgsStr = tgsStr s
                 <> T.pack (show i) <> " [shape=egg,color=green,label=\"In:" <> showsz (tvTensorSizes v) <> "\"]" <> "\n"
        , tgsMap = M.insert (tvTensorName v) i (tgsMap s) }
    mkNode n s i | any (\x -> T.isPrefixOf (T.pack x) (tnKind n) || T.isSuffixOf (T.pack x) (tnKind n)) censoredNodes = s
                 | noConstants && tnIsConstant n = s
                 | otherwise =
      s { tgsStr = tgsStr s
                 <> T.pack (show i) <> " [shape=box,"
                                    <> (if tnIsConstant n then
                                          "color=orange," else
                                          "")
                                    <> "label=\""
                                    <> (if tnIsConstant n then
                                         "Const:" <> (T.intercalate " "
                                                     (map showsz (map tvTensorSizes (tnOutputs n)))) else
                                       tnKind n)
                                  <>"\"]" <> "\n"
                 <> T.unlines
                       (map (\x ->
                               case M.lookup (tvTensorName x) (tgsMap s) of
                                 Nothing ->
                                   if noConstants then
                                     "" else
                                     error $ "Can't find an input? This is a bug! " <> show x <> show n
                                 Just x -> T.pack (show x) <> " -> "  <> T.pack (show i))
                            (tnInputs n))
        , tgsMap = L.foldl' (\m v -> M.insert (tvTensorName v) i m) (tgsMap s) (tnOutputs n)
        , tgsOutputs =
            (L.foldl' (\m v ->
                       case M.lookup v m of
                         Just x  -> M.adjust (i:) v m
                         Nothing -> m)
              (L.foldl' (flip M.delete) (tgsOutputs s) (tnInputs n))
              (tnOutputs n))
        }
    mkOutputNode v is =
         "O" <> (tvTensorName v) <> " [shape=hexagon,color=blue,label=\"Out:" <> showsz (tvTensorSizes v) <> "\"]" <> "\n"
      <> T.unlines (map (\i -> T.pack (show i) <> " -> O" <> (tvTensorName v)) is)
    loopGraph state g =
      let state' = L.foldl' (\s v ->
                            case M.lookup (tvTensorName v) (tgsMap s) of
                              Nothing -> withNextNode s (mkInputNode v)
                              Just _  -> s)
                          state
                          (traceGraphInputs trace)
      in let state'' = L.foldl' (\s v -> s { tgsOutputs = M.insert v [] (tgsOutputs s) } )
                                state'
                                (traceGraphOutputs trace)
         in L.foldl' loopNode state'' (traceGraphNodes g)
    loopNode s n = withNextNode s (mkNode n)
    censoredNodes = ["aten::_convolution_nogroup"
                    ,"_forward"
                    ,"aten::thnn_"]
