masahi commented on a change in pull request #8599:
URL: https://github.com/apache/tvm/pull/8599#discussion_r682191113
##########
File path: python/tvm/relay/frontend/common.py
##########
@@ -624,3 +624,128 @@ def to_int_list(np_array):
cause problems in relay/TOPI.
"""
return [int(x) for x in np_array]
+
+
+def unbind(data, axis=0):
+ """
+ Unbind was gotten from pytorch.py and modified. The operation removes a
tensor dimension
+ and returns a tuple of all slices along a given dimension, already without
it.
+ TODO (vvchernov): It needs such operation on relay side to reduce time
consumption
+ on squeeze operation.
+
+ Parameters
+ ----------
+ data : relay.Expr
+ Input tensor
+ axis : int
+ Axis along which tensor is splited. Tensors in the list do not have
this axis.
+ Returns
+ -------
+ result : List[relay.Expr]
+ The sequence of computed tensors
+ """
+ shape = infer_shape(data)
+ if axis >= len(shape):
+ msg = "Please check input dim, it shouldn't be greater than or equal
to rank."
+ raise AttributeError(msg)
+
+ selections = shape[axis]
+ res_split = _op.split(data, selections, axis)
+ ret = []
+ for i in range(selections):
+ ret.append(_op.squeeze(res_split[i], axis=[axis]))
+ return _expr.TupleWrapper(_expr.Tuple(ret), selections)
+
+
+def lstm_cell(
+ input_seqs,
+ hidden_state,
+ cell_state,
+ w_inp,
+ w_hid,
+ b_inp=None,
+ b_hid=None,
+ proj=None,
+ p_i=None,
+ p_f=None,
+ p_o=None,
+ f_act=_op.sigmoid,
+ g_act=_op.tanh,
+ h_act=_op.tanh,
+ backwards=False,
+):
+ """
+ Common implementation of LSTM cell for all frontends of TVM
+ TODO (vvchernov): currently it is used by onnx and pytorch. Extend for
other frontends
+
+ Parameters
+ ----------
+ input_seqs : List[relay.Expr]
+ The sequence of input tensors
+ Input tensor should be 2d while issue #8412 is not resolved
+ Shape = (batch, feature_size)
+ hidden_state : relay.Expr
+ Hidden state. shape = (batch, hidden_size)
+ cell_state : relay.Expr
+ Cell state. shape = (batch, hidden_size)
+ w_inp, w_hid : relay.Expr
+ weight matrices. wi shape = (4 * hidden_size, feature_size)
+ wh shape = (4 * hidden_size, hidden_size or proj_size)
+ NOTE: wi = (w_ii|w_if|w_ig|w_io) for input, forget, cell and output
gates.
+ The order is important for correct LSTM calculation!
+ b_inp, b_hid : relay.Expr
+ bias matrices. The same order of internal parts as for weights. shape
= (4 * hidden_size)
+ proj : relay.Expr
+ projection matrix. shape = (proj_size, hidden_size)
+ p_i, p_f, p_o : relay.Expr
+ peephole LSTM matrices. shape = (batch, hidden_size)
+ f_act, g_act, h_act : relay.op
+ activation funtions
+ backwards : bool
+ Flag for reverse pass of LSTM
+
+ Returns
+ -------
+ result : List[relay.Expr], relay.Expr, relay.Expr
+ The sequence of computed result, final hidden and cell state
+ """
+
+ outputs_list = []
+ for x_t in input_seqs if not backwards else reversed(input_seqs):
+ # x_t shape = (batch, feature size), step shape = (batch, feature size
+ hidden_size)
+ step = _op.concatenate([x_t, hidden_state], axis=1)
+ cat_w = _op.concatenate([w_inp, w_hid], axis=1)
+ # Instead of nn.dense(x_t, w_inp) + nn.dense(hidden_state, w_hid)
+ # the nn.dense(step, cat_w) is used
Review comment:
remove the
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
