code: def lstm(
      Xs,
      Wi,
      Wh,
      Bi=None,
      Bh=None,
      h_init=None,
      c_init=None,
      proj=None,
      p_i=None,
      p_f=None,
      p_o=None,
      f_act=tir.sigmoid,
      g_act=tir.tanh,
      h_act=tir.tanh,
      reverse=False,
      weight_layout: str = "IFGO",
  ):
      """General LSTM implemented using TE scan.
  
      Parameters
      ----------
      Xs : te.Tensor
          Input sequence with shape `(seq_len, batch_size, in_dim)`
      Wi : te.Tensor
          Input weight matrix with shape `(4 * hidden_dim, in_dim)`. The weights are packed according
          to `weight_layout`.
      Wh : te.Tensor
          Hidden weight matrix with shape `(4 * hidden_dim, hidden_dim or proj_dim)`. Packed as `Wh`.
      Bi : te.Tensor, optional
          Input bias with shape `(4 * hidden_dim,)`, by default None. Packed as `Wh`.
      Bh : te.Tensor, optional
          Hidden bias with shape as `Bi`, by default None. Packed as `Wh`.
      h_init : te.Tensor, optional
          Initial hidden state with shape `(batch_size, hidden_dim or proj_dim)`, zero if None
      c_init : te.Tensor, optional
          Initial cell state with same shape as `h_init`, zero if None
      proj : te.Tensor, optional
          Projection matrix with shape `(proj_dim, hidden_dim)`, by default None
      p_i, p_f, p_o : te.Tensor, optional
          Peephole LSTM matrices with shape `(batch_size, hidden_dim)`, by default None
      f_act, g_act, h_act : F, optional
          Gate activation functions
      reverse : bool, optional
          Whether to process `Xs` in reverse, by default False
      weight_layout : str, optional
          The packed weight layout for gates, by default "IFGO". Note: I = input, F = forget,
          G = cell, O = output.
  
      Returns
      -------
      result : te.Tensor, te.Tensor
          Tuple of hidden states (with shape `(seq_len, batch_size, hidden_dim or proj_dim)`), and
          cell states (with shape `(seq_len, batch_size, hidden_dim)`).
      """
      assert len(weight_layout) == 4 and sorted(
          weight_layout
      ) == sorted(
          "IFGO"
      ), f'given weight layout "{weight_layout}" is not a permutation of "IFGO"'
  
      i_gate_idx = weight_layout.find("I")
      f_gate_idx = weight_layout.find("F")
      g_gate_idx = weight_layout.find("G")
      o_gate_idx = weight_layout.find("O")
  
      seq_len, batch_size, in_dim = Xs.shape
      assert (
          Wi.shape[0] % 4 == 0
      ), f"dim 0 of input weight should be 4 * hidden_dim, but {Wi.shape[0]} is not divisible by 4"
      hidden_dim = Wi.shape[0] // 4
      proj_dim = hidden_dim
      if proj is not None:
          proj_dim = proj.shape[0]
  
      # te.scan uses up 1 element for the initial value
      scan_len = seq_len + 1
  
      # precompute input-hidden matmul outside the scan
      ki = te.reduce_axis((0, in_dim), name="ki2h")
      Xi2h = te.compute(
          (seq_len * batch_size, 4 * hidden_dim),
          lambda tb, ij: te.sum(
              Xs[(tb // batch_size), tb % batch_size, ki] * Wi[ij, ki],
              axis=ki,
          ),
          name="Xi2h",
      )
      if Bi is not None:
          Xi2h = te.compute(
              Xi2h.shape,
              lambda tb, ij: Xi2h[tb, ij] + Bi[ij],
              name="Xi2h_bias",
              tag=tag.INJECTIVE,
          )
  
      h_state = te.placeholder(
          (scan_len, batch_size, proj_dim), name="h_state"
      )
      c_state = te.placeholder(
          (scan_len, batch_size, hidden_dim), name="c_state"
      )
      h_init = te.compute(
          (1, batch_size, proj_dim),
          lambda _, b, i: h_init[b, i] if h_init is not None else 0.0,
          name="h_init",
      )
      c_init = te.compute(
          (1, batch_size, hidden_dim),
          lambda _, b, i: c_init[b, i] if c_init is not None else 0.0,
          name="c_init",
      )
  
      # begin scan computations, first the (batched) hidden-hidden dense
      kh = te.reduce_axis((0, proj_dim), name="kh2h")
      s_h2h = te.compute(
          (scan_len, batch_size, 4, hidden_dim),
          lambda t, b, i, j: te.sum(
              h_state[t - 1, b, kh] * Wh[i * hidden_dim + j, kh],
              axis=kh,
          ),
          name="s_h2h",
      )
      if Bh is not None:
          s_h2h = te.compute(
              s_h2h.shape,
              lambda t, b, i, j: s_h2h[t, b, i, j]
              + Bh[i * hidden_dim + j],
              name="s_h2h_bias",
              tag=tag.INJECTIVE,
          )
  
      # helper to reverse time if scanning backwards
      get_x_t = lambda t: seq_len - t if reverse else t - 1
  
      gates = te.compute(
          (scan_len, batch_size, 4, hidden_dim),
          lambda t, b, i, j: Xi2h[
              get_x_t(t) * batch_size + b, i * hidden_dim + j
          ]
          + s_h2h[t, b, i, j],
          name="gates",
          tag=tag.INJECTIVE,
      )
  
      # helper to correctly read each gate dense from the batched output
      read_gate = lambda t, b, j, idx: gates[t, b, idx, j]
  
      gate_shape = (scan_len, batch_size, hidden_dim)
  
      # compute the activated gates (and do some extra stuff if peephole weights are present)
      if p_i is not None and p_f is not None:
          i_gate = te.compute(
              gate_shape,
              lambda t, b, j: f_act(
                  read_gate(t, b, j, i_gate_idx)
                  + p_i[b, j] * c_state[t - 1, b, j]
              ),
              name="i_gate_p",
              tag=tag.INJECTIVE,
          )
          f_gate = te.compute(
              gate_shape,
              lambda t, b, j: f_act(
                  read_gate(t, b, j, f_gate_idx)
                  + p_f[b, j] * c_state[t - 1, b, j]
              ),
              name="f_gate_p",
              tag=tag.INJECTIVE,
          )
      else:
          i_gate = te.compute(
              gate_shape,
              lambda *i: f_act(read_gate(*i, i_gate_idx)),
              name="i_gate",
              tag=tag.INJECTIVE,
          )
          f_gate = te.compute(
              gate_shape,
              lambda *i: f_act(read_gate(*i, f_gate_idx)),
              name="f_gate",
              tag=tag.INJECTIVE,
          )
  
      g_gate = te.compute(
          gate_shape,
          lambda *i: g_act(read_gate(*i, g_gate_idx)),
          name="g_gate",
          tag=tag.INJECTIVE,
      )
  
      next_c = te.compute(
          gate_shape,
          lambda t, b, j: f_gate[t, b, j] * c_state[t - 1, b, j]
          + i_gate[t, b, j] * g_gate[t, b, j],
          name="next_c",
      )
  
      if p_o is not None:
          o_gate = te.compute(
              gate_shape,
              lambda t, b, j: f_act(
                  read_gate(t, b, j, o_gate_idx)
                  + p_o[b, j] * next_c[t, b, j]
              ),
              name="o_gate_p",
              tag=tag.INJECTIVE,
          )
      else:
          o_gate = te.compute(
              gate_shape,
              lambda *i: f_act(read_gate(*i, o_gate_idx)),
              name="o_gate",
              tag=tag.INJECTIVE,
          )
  
      next_h = te.compute(
          gate_shape,
          lambda *i: o_gate(*i) * h_act(next_c(*i)),
          name="next_h",
      )
  
      # project hidden state back to proj_dim if projection matrix is present
      if proj is not None:
          kr = te.reduce_axis((0, hidden_dim), name="kh2p")
          next_h = te.compute(
              (scan_len, batch_size, proj_dim),
              lambda t, b, j: te.sum(
                  next_h[t, b, kr] * proj[j, kr], axis=kr
              ),
              name="next_h_proj",
          )
  
      scan_h, scan_c = te.scan(
          [h_init, c_init],
          [next_h, next_c],
          [h_state, c_state],
          name="lstm_scan",
      )
  
      # drop the initial values, TODO(@altanh): is there a better way?
      scan_h = te.compute(
          (seq_len, batch_size, proj_dim),
          lambda t, b, j: scan_h[t + 1, b, j],
          name="hidden_states",
      )
      scan_c = te.compute(
          (seq_len, batch_size, hidden_dim),
          lambda t, b, j: scan_c[t + 1, b, j],
          name="cell_states",
      )
  
      return scan_h, scan_c
  
label_desc: Use image node_id 1 for calling a global function with example usage: lstm(Xs, Wi, Wh, Bi, Bh, h_init, c_init, proj, p_i, p_f, p_o, f_act, g_act, h_act, reverse, weight_layout) and returns: scan_h,  scan_c