rnn_ir.txt

graph(%self : __torch__.custom_lstms.LSTMLayer,
      %input.1 : Tensor,
      %state.1 : (Tensor, Tensor)):
  %3 : bool = prim::Constant[value=1]() # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:84:8
  %4 : int = prim::Constant[value=0]() # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:84:34
  %outputs.1 : Tensor[] = prim::ListConstruct()
  %6 : int = aten::size(%input.1, %4) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:84:23
  %outputs : Tensor[], %state : (Tensor, Tensor) = prim::Loop(%6, %3, %outputs.1, %state.1) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:84:8
    block0(%i.1 : int, %outputs.6 : Tensor[], %state.6 : (Tensor, Tensor)):
      %12 : __torch__.custom_lstms.LSTMCell = prim::GetAttr[name="cell"](%self)
      %13 : Tensor = aten::select(%input.1, %4, %i.1) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:85:35
      %21 : int = prim::Constant[value=4]() # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:27:60
      %22 : int = prim::Constant[value=1]() # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:27:63
      %hx.1 : Tensor, %cx.1 : Tensor = prim::TupleUnpack(%state.6)
      %25 : Tensor = prim::GetAttr[name="weight_ih"](%12)
      %26 : Tensor = aten::t(%25) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:25:33
      %27 : Tensor = aten::mm(%13, %26) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:25:17
      %28 : Tensor = prim::GetAttr[name="bias_ih"](%12)
      %29 : Tensor = aten::add(%27, %28, %22) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:25:17
      %30 : Tensor = prim::GetAttr[name="weight_hh"](%12)
      %31 : Tensor = aten::t(%30) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:26:30
      %32 : Tensor = aten::mm(%hx.1, %31) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:26:17
      %33 : Tensor = aten::add(%29, %32, %22) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:25:17
      %34 : Tensor = prim::GetAttr[name="bias_hh"](%12)
      %gates.1 : Tensor = aten::add(%33, %34, %22) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:25:17
      %36 : Tensor[] = aten::chunk(%gates.1, %21, %22) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:27:48
      %ingate.1 : Tensor, %forgetgate.1 : Tensor, %cellgate.1 : Tensor, %outgate.1 : Tensor = prim::ListUnpack(%36)
      %ingate.3 : Tensor = aten::sigmoid(%ingate.1) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:29:17
      %forgetgate.3 : Tensor = aten::sigmoid(%forgetgate.1) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:30:21
      %cellgate.3 : Tensor = aten::tanh(%cellgate.1) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:31:19
      %outgate.3 : Tensor = aten::sigmoid(%outgate.1) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:32:18
      %45 : Tensor = aten::mul(%forgetgate.3, %cx.1) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:34:14
      %46 : Tensor = aten::mul(%ingate.3, %cellgate.3) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:34:34
      %cy.1 : Tensor = aten::add(%45, %46, %22) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:34:14
      %48 : Tensor = aten::tanh(%cy.1) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:35:23
      %hy.1 : Tensor = aten::mul(%outgate.3, %48) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:35:13
      %50 : (Tensor, Tensor) = prim::TupleConstruct(%hy.1, %cy.1)
      %51 : (Tensor, (Tensor, Tensor)) = prim::TupleConstruct(%hy.1, %50)
      %out.1 : Tensor, %state.3 : (Tensor, Tensor) = prim::TupleUnpack(%51)
      %17 : Tensor[] = prim::ListConstruct(%out.1)
      %outputs.3 : Tensor[] = aten::add_(%outputs.6, %17) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:86:12
      -> (%3, %outputs.3, %state.3)
  %19 : Tensor = aten::stack(%outputs, %4) # /home/masa/projects/dev/torchscript-to-tvm/custom_lstms.py:87:15
  %20 : (Tensor, (Tensor, Tensor)) = prim::TupleConstruct(%19, %state)
  return (%20)

['input.1', 'state.1']
[('X', (5, 2, 3)), ('states', ((2, 4), (2, 4)))]