gemini-code-assist[bot] commented on code in PR #19633:
URL: https://github.com/apache/tvm/pull/19633#discussion_r3315244158
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -4878,6 +4880,222 @@ def convert_unpack(self, op):
return squeezed
+ def convert_lstm(self, op):
+ """Convert TFLite LSTM (single-step).
+
+ Standard LSTM cell with FULL kernel and coupled input-forget gate.
+ Peephole, projection, and layer norm are not supported.
+
+ Inputs (24 tensors, many optional):
+ [0] input [batch, input_size]
+ [1] input_to_input_weights (optional, -1 => coupled)
+ [2] input_to_forget_weights [num_units, input_size]
+ [3] input_to_cell_weights [num_units, input_size]
+ [4] input_to_output_weights [num_units, input_size]
+ [5] recurrent_to_input_weights (optional)
+ [6] recurrent_to_forget_weights [num_units, num_units]
+ [7] recurrent_to_cell_weights [num_units, num_units]
+ [8] recurrent_to_output_weights [num_units, num_units]
+ [9-11] cell_to_*_weights (optional, not supported)
+ [12] input_gate_bias (optional)
+ [13] forget_gate_bias [num_units]
+ [14] cell_bias [num_units]
+ [15] output_gate_bias [num_units]
+ [16-17] projection_weights/bias (optional, not supported)
+ [18] output_state [batch, num_units]
+ [19] cell_state [batch, num_units]
+ [20-23] layer_norm (optional, not supported)
+
+ Output:
+ [0] output [batch, num_units]
+
+ Cell (coupled input-forget):
+ f = sigmoid(x @ W_f.T + h @ R_f.T + b_f)
+ i = 1 - f
+ g = tanh(x @ W_c.T + h @ R_c.T + b_c)
+ o = sigmoid(x @ W_o.T + h @ R_o.T + b_o)
+ c_new = f * c_prev + i * g
+ h_new = fused_activation(o * tanh(c_new))
+ """
+ from tflite.BuiltinOptions import BuiltinOptions
+ from tflite.LSTMOptions import LSTMOptions
+
+ if self.is_quantized(op):
+ raise tvm.error.OpNotImplemented(
+ "TFLite quantized LSTM is not supported yet."
+ )
+
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 24, (
+ f"input tensors length should be 24, got {len(input_tensors)}"
+ )
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) >= 1, "output tensors length should be at
least 1"
+
+ assert op.BuiltinOptionsType() == BuiltinOptions.LSTMOptions
+ op_options = op.BuiltinOptions()
+ lstm_opts = LSTMOptions()
+ lstm_opts.Init(op_options.Bytes, op_options.Pos)
+
+ fused_activation_fn = lstm_opts.FusedActivationFunction()
+ cell_clip = lstm_opts.CellClip()
+ proj_clip = lstm_opts.ProjClip()
+
+ in_expr = self.get_tensor_expr(input_tensors[0])
+
+ # Only coupled input-forget gate is supported.
+ assert input_tensors[1].tensor_idx == -1, (
+ "Only coupled input-forget LSTM is supported."
+ )
+ assert input_tensors[5].tensor_idx == -1, (
+ "Only coupled input-forget LSTM is supported."
+ )
+
+ # Weights.
+ w_f = self.get_tensor_expr(input_tensors[2])
+ w_c = self.get_tensor_expr(input_tensors[3])
+ w_o = self.get_tensor_expr(input_tensors[4])
+
+ r_f = self.get_tensor_expr(input_tensors[6])
+ r_c = self.get_tensor_expr(input_tensors[7])
+ r_o = self.get_tensor_expr(input_tensors[8])
+
+ # Biases.
+ b_f = self.get_tensor_expr(input_tensors[13])
+ b_c = self.get_tensor_expr(input_tensors[14])
+ b_o = self.get_tensor_expr(input_tensors[15])
+
+ # State inputs.
+ h_prev = self.get_tensor_expr(input_tensors[18])
+ c_prev = self.get_tensor_expr(input_tensors[19])
+
+ # Coupled input-forget gate.
+ f = relax.op.sigmoid(
+ relax.op.add(
+ relax.op.add(
+ relax.op.matmul(in_expr, relax.op.permute_dims(w_f)),
+ relax.op.matmul(h_prev, relax.op.permute_dims(r_f)),
+ ),
+ b_f,
+ )
+ )
+ i = relax.op.subtract(
+ relax.const(np.array(1.0, dtype="float32"), "float32"),
+ f,
+ )
+
+ # Cell candidate.
+ g = relax.op.tanh(
+ relax.op.add(
+ relax.op.add(
+ relax.op.matmul(in_expr, relax.op.permute_dims(w_c)),
+ relax.op.matmul(h_prev, relax.op.permute_dims(r_c)),
+ ),
+ b_c,
+ )
+ )
+
+ # Output gate.
+ o = relax.op.sigmoid(
+ relax.op.add(
+ relax.op.add(
+ relax.op.matmul(in_expr, relax.op.permute_dims(w_o)),
+ relax.op.matmul(h_prev, relax.op.permute_dims(r_o)),
+ ),
+ b_o,
+ )
+ )
+
+ # Cell state update with optional clipping.
+ c_new = relax.op.add(
+ relax.op.multiply(f, c_prev),
+ relax.op.multiply(i, g),
+ )
+ if cell_clip > 0:
+ c_new = relax.op.clip(c_new, -cell_clip, cell_clip)
+
+ # Hidden state.
+ h_new = relax.op.multiply(o, relax.op.tanh(c_new))
+ if proj_clip > 0:
+ h_new = relax.op.clip(h_new, -proj_clip, proj_clip)
+
+ return self.convert_fused_activation_function(h_new,
fused_activation_fn)
+
+ def convert_svdf(self, op):
+ """Convert TFLite SVDF (single-step).
+
+ Structured-Vectorized Bidirectional Filter for keyword spotting.
+
+ Inputs (5 tensors):
+ [0] input [batch, input_size]
+ [1] feature_weights [num_filters, input_size]
+ [2] time_weights [num_filters, rank]
+ [3] bias [num_filters] (optional)
+ [4] state [batch, num_filters * rank] (variable)
+
+ Output:
+ [0] output [batch, num_filters]
+
+ Computation:
+ feat = x @ W_feat.T # feature projection
+ state_r = reshape(state, [B, F, rank]) # reshape
+ time = sum(state_r * time_weights, axis=-1) # time filtering
+ out = activation(feat + time + bias)
+ """
+ from tflite.BuiltinOptions import BuiltinOptions
+ from tflite.SVDFOptions import SVDFOptions
+
+ if self.is_quantized(op):
+ raise tvm.error.OpNotImplemented(
+ "TFLite quantized SVDF is not supported yet."
+ )
+
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 5, (
+ f"input tensors length should be 5, got {len(input_tensors)}"
+ )
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) >= 1, "output tensors length should be at
least 1"
+
+ assert op.BuiltinOptionsType() == BuiltinOptions.SVDFOptions
+ op_options = op.BuiltinOptions()
+ svdf_opts = SVDFOptions()
+ svdf_opts.Init(op_options.Bytes, op_options.Pos)
+
+ rank = svdf_opts.Rank()
+ fused_activation_fn = svdf_opts.FusedActivationFunction()
+
+ in_expr = self.get_tensor_expr(input_tensors[0])
+ feat_weights = self.get_tensor_expr(input_tensors[1])
+ time_weights = self.get_tensor_expr(input_tensors[2])
+
+ batch_size = to_int_list(self.get_tensor_shape(input_tensors[0]))[0]
Review Comment:
Using `to_int_list` on the input tensor shape forces the batch dimension to
be a static Python integer. If the model has a dynamic batch size (represented
as a symbolic variable like `tir.Var`), `to_int_list` will raise a `TypeError`
and crash during model import. Extracting the batch size directly and only
converting to `int` if it is a static integer preserves symbolic dimensions for
dynamic batching.
```suggestion
batch_size = self.get_tensor_shape(input_tensors[0])[0]
if isinstance(batch_size, (np.integer, int)):
batch_size = int(batch_size)
```
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -4878,6 +4880,222 @@ def convert_unpack(self, op):
return squeezed
+ def convert_lstm(self, op):
+ """Convert TFLite LSTM (single-step).
+
+ Standard LSTM cell with FULL kernel and coupled input-forget gate.
+ Peephole, projection, and layer norm are not supported.
+
+ Inputs (24 tensors, many optional):
+ [0] input [batch, input_size]
+ [1] input_to_input_weights (optional, -1 => coupled)
+ [2] input_to_forget_weights [num_units, input_size]
+ [3] input_to_cell_weights [num_units, input_size]
+ [4] input_to_output_weights [num_units, input_size]
+ [5] recurrent_to_input_weights (optional)
+ [6] recurrent_to_forget_weights [num_units, num_units]
+ [7] recurrent_to_cell_weights [num_units, num_units]
+ [8] recurrent_to_output_weights [num_units, num_units]
+ [9-11] cell_to_*_weights (optional, not supported)
+ [12] input_gate_bias (optional)
+ [13] forget_gate_bias [num_units]
+ [14] cell_bias [num_units]
+ [15] output_gate_bias [num_units]
+ [16-17] projection_weights/bias (optional, not supported)
+ [18] output_state [batch, num_units]
+ [19] cell_state [batch, num_units]
+ [20-23] layer_norm (optional, not supported)
+
+ Output:
+ [0] output [batch, num_units]
+
+ Cell (coupled input-forget):
+ f = sigmoid(x @ W_f.T + h @ R_f.T + b_f)
+ i = 1 - f
+ g = tanh(x @ W_c.T + h @ R_c.T + b_c)
+ o = sigmoid(x @ W_o.T + h @ R_o.T + b_o)
+ c_new = f * c_prev + i * g
+ h_new = fused_activation(o * tanh(c_new))
+ """
+ from tflite.BuiltinOptions import BuiltinOptions
+ from tflite.LSTMOptions import LSTMOptions
+
+ if self.is_quantized(op):
+ raise tvm.error.OpNotImplemented(
+ "TFLite quantized LSTM is not supported yet."
+ )
+
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 24, (
+ f"input tensors length should be 24, got {len(input_tensors)}"
+ )
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) >= 1, "output tensors length should be at
least 1"
+
+ assert op.BuiltinOptionsType() == BuiltinOptions.LSTMOptions
+ op_options = op.BuiltinOptions()
+ lstm_opts = LSTMOptions()
+ lstm_opts.Init(op_options.Bytes, op_options.Pos)
+
+ fused_activation_fn = lstm_opts.FusedActivationFunction()
+ cell_clip = lstm_opts.CellClip()
+ proj_clip = lstm_opts.ProjClip()
+
+ in_expr = self.get_tensor_expr(input_tensors[0])
+
+ # Only coupled input-forget gate is supported.
+ assert input_tensors[1].tensor_idx == -1, (
+ "Only coupled input-forget LSTM is supported."
+ )
+ assert input_tensors[5].tensor_idx == -1, (
+ "Only coupled input-forget LSTM is supported."
+ )
+
+ # Weights.
+ w_f = self.get_tensor_expr(input_tensors[2])
+ w_c = self.get_tensor_expr(input_tensors[3])
+ w_o = self.get_tensor_expr(input_tensors[4])
+
+ r_f = self.get_tensor_expr(input_tensors[6])
+ r_c = self.get_tensor_expr(input_tensors[7])
+ r_o = self.get_tensor_expr(input_tensors[8])
+
+ # Biases.
+ b_f = self.get_tensor_expr(input_tensors[13])
+ b_c = self.get_tensor_expr(input_tensors[14])
+ b_o = self.get_tensor_expr(input_tensors[15])
+
+ # State inputs.
+ h_prev = self.get_tensor_expr(input_tensors[18])
+ c_prev = self.get_tensor_expr(input_tensors[19])
+
+ # Coupled input-forget gate.
+ f = relax.op.sigmoid(
+ relax.op.add(
+ relax.op.add(
+ relax.op.matmul(in_expr, relax.op.permute_dims(w_f)),
+ relax.op.matmul(h_prev, relax.op.permute_dims(r_f)),
+ ),
+ b_f,
+ )
+ )
+ i = relax.op.subtract(
+ relax.const(np.array(1.0, dtype="float32"), "float32"),
+ f,
+ )
Review Comment:
The constant `1.0` can be created directly using `relax.const(1.0,
"float32")` instead of wrapping a NumPy array. This is cleaner, more readable,
and avoids unnecessary NumPy array creation overhead.
```suggestion
i = relax.op.subtract(
relax.const(1.0, "float32"),
f,
)
```
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -4878,6 +4880,222 @@ def convert_unpack(self, op):
return squeezed
+ def convert_lstm(self, op):
+ """Convert TFLite LSTM (single-step).
+
+ Standard LSTM cell with FULL kernel and coupled input-forget gate.
+ Peephole, projection, and layer norm are not supported.
+
+ Inputs (24 tensors, many optional):
+ [0] input [batch, input_size]
+ [1] input_to_input_weights (optional, -1 => coupled)
+ [2] input_to_forget_weights [num_units, input_size]
+ [3] input_to_cell_weights [num_units, input_size]
+ [4] input_to_output_weights [num_units, input_size]
+ [5] recurrent_to_input_weights (optional)
+ [6] recurrent_to_forget_weights [num_units, num_units]
+ [7] recurrent_to_cell_weights [num_units, num_units]
+ [8] recurrent_to_output_weights [num_units, num_units]
+ [9-11] cell_to_*_weights (optional, not supported)
+ [12] input_gate_bias (optional)
+ [13] forget_gate_bias [num_units]
+ [14] cell_bias [num_units]
+ [15] output_gate_bias [num_units]
+ [16-17] projection_weights/bias (optional, not supported)
+ [18] output_state [batch, num_units]
+ [19] cell_state [batch, num_units]
+ [20-23] layer_norm (optional, not supported)
+
+ Output:
+ [0] output [batch, num_units]
+
+ Cell (coupled input-forget):
+ f = sigmoid(x @ W_f.T + h @ R_f.T + b_f)
+ i = 1 - f
+ g = tanh(x @ W_c.T + h @ R_c.T + b_c)
+ o = sigmoid(x @ W_o.T + h @ R_o.T + b_o)
+ c_new = f * c_prev + i * g
+ h_new = fused_activation(o * tanh(c_new))
+ """
+ from tflite.BuiltinOptions import BuiltinOptions
+ from tflite.LSTMOptions import LSTMOptions
+
+ if self.is_quantized(op):
+ raise tvm.error.OpNotImplemented(
+ "TFLite quantized LSTM is not supported yet."
+ )
+
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 24, (
+ f"input tensors length should be 24, got {len(input_tensors)}"
+ )
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) >= 1, "output tensors length should be at
least 1"
+
+ assert op.BuiltinOptionsType() == BuiltinOptions.LSTMOptions
+ op_options = op.BuiltinOptions()
+ lstm_opts = LSTMOptions()
+ lstm_opts.Init(op_options.Bytes, op_options.Pos)
+
+ fused_activation_fn = lstm_opts.FusedActivationFunction()
+ cell_clip = lstm_opts.CellClip()
+ proj_clip = lstm_opts.ProjClip()
+
+ in_expr = self.get_tensor_expr(input_tensors[0])
+
+ # Only coupled input-forget gate is supported.
+ assert input_tensors[1].tensor_idx == -1, (
+ "Only coupled input-forget LSTM is supported."
+ )
+ assert input_tensors[5].tensor_idx == -1, (
+ "Only coupled input-forget LSTM is supported."
+ )
+
+ # Weights.
+ w_f = self.get_tensor_expr(input_tensors[2])
+ w_c = self.get_tensor_expr(input_tensors[3])
+ w_o = self.get_tensor_expr(input_tensors[4])
+
+ r_f = self.get_tensor_expr(input_tensors[6])
+ r_c = self.get_tensor_expr(input_tensors[7])
+ r_o = self.get_tensor_expr(input_tensors[8])
+
+ # Biases.
+ b_f = self.get_tensor_expr(input_tensors[13])
+ b_c = self.get_tensor_expr(input_tensors[14])
+ b_o = self.get_tensor_expr(input_tensors[15])
+
+ # State inputs.
+ h_prev = self.get_tensor_expr(input_tensors[18])
+ c_prev = self.get_tensor_expr(input_tensors[19])
+
+ # Coupled input-forget gate.
+ f = relax.op.sigmoid(
+ relax.op.add(
+ relax.op.add(
+ relax.op.matmul(in_expr, relax.op.permute_dims(w_f)),
+ relax.op.matmul(h_prev, relax.op.permute_dims(r_f)),
+ ),
+ b_f,
+ )
+ )
+ i = relax.op.subtract(
+ relax.const(np.array(1.0, dtype="float32"), "float32"),
+ f,
+ )
+
+ # Cell candidate.
+ g = relax.op.tanh(
+ relax.op.add(
+ relax.op.add(
+ relax.op.matmul(in_expr, relax.op.permute_dims(w_c)),
+ relax.op.matmul(h_prev, relax.op.permute_dims(r_c)),
+ ),
+ b_c,
+ )
+ )
+
+ # Output gate.
+ o = relax.op.sigmoid(
+ relax.op.add(
+ relax.op.add(
+ relax.op.matmul(in_expr, relax.op.permute_dims(w_o)),
+ relax.op.matmul(h_prev, relax.op.permute_dims(r_o)),
+ ),
+ b_o,
+ )
+ )
+
+ # Cell state update with optional clipping.
+ c_new = relax.op.add(
+ relax.op.multiply(f, c_prev),
+ relax.op.multiply(i, g),
+ )
+ if cell_clip > 0:
+ c_new = relax.op.clip(c_new, -cell_clip, cell_clip)
+
+ # Hidden state.
+ h_new = relax.op.multiply(o, relax.op.tanh(c_new))
+ if proj_clip > 0:
+ h_new = relax.op.clip(h_new, -proj_clip, proj_clip)
+
+ return self.convert_fused_activation_function(h_new,
fused_activation_fn)
Review Comment:
In TFLite, LSTM states (`h_prev` and `c_prev`) are mutable variables that
are updated at each step. In TVM Relax, which is functional, we should update
the state expressions in the `ExprTable` using `force_override=True` so that
subsequent operations or steps in the graph correctly receive the updated
states.
```python
# Hidden state.
h_new = relax.op.multiply(o, relax.op.tanh(c_new))
if proj_clip > 0:
h_new = relax.op.clip(h_new, -proj_clip, proj_clip)
h_activated = self.convert_fused_activation_function(h_new,
fused_activation_fn)
# Update state tensors in the expression table for subsequent ops
self.exp_tab.set_expr(
get_tensor_name(self.subgraph, input_tensors[18].tensor_idx),
h_activated,
force_override=True,
)
self.exp_tab.set_expr(
get_tensor_name(self.subgraph, input_tensors[19].tensor_idx),
c_new,
force_override=True,
)
return h_activated
```
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -4878,6 +4880,222 @@ def convert_unpack(self, op):
return squeezed
+ def convert_lstm(self, op):
+ """Convert TFLite LSTM (single-step).
+
+ Standard LSTM cell with FULL kernel and coupled input-forget gate.
+ Peephole, projection, and layer norm are not supported.
+
+ Inputs (24 tensors, many optional):
+ [0] input [batch, input_size]
+ [1] input_to_input_weights (optional, -1 => coupled)
+ [2] input_to_forget_weights [num_units, input_size]
+ [3] input_to_cell_weights [num_units, input_size]
+ [4] input_to_output_weights [num_units, input_size]
+ [5] recurrent_to_input_weights (optional)
+ [6] recurrent_to_forget_weights [num_units, num_units]
+ [7] recurrent_to_cell_weights [num_units, num_units]
+ [8] recurrent_to_output_weights [num_units, num_units]
+ [9-11] cell_to_*_weights (optional, not supported)
+ [12] input_gate_bias (optional)
+ [13] forget_gate_bias [num_units]
+ [14] cell_bias [num_units]
+ [15] output_gate_bias [num_units]
+ [16-17] projection_weights/bias (optional, not supported)
+ [18] output_state [batch, num_units]
+ [19] cell_state [batch, num_units]
+ [20-23] layer_norm (optional, not supported)
+
+ Output:
+ [0] output [batch, num_units]
+
+ Cell (coupled input-forget):
+ f = sigmoid(x @ W_f.T + h @ R_f.T + b_f)
+ i = 1 - f
+ g = tanh(x @ W_c.T + h @ R_c.T + b_c)
+ o = sigmoid(x @ W_o.T + h @ R_o.T + b_o)
+ c_new = f * c_prev + i * g
+ h_new = fused_activation(o * tanh(c_new))
+ """
+ from tflite.BuiltinOptions import BuiltinOptions
+ from tflite.LSTMOptions import LSTMOptions
+
+ if self.is_quantized(op):
+ raise tvm.error.OpNotImplemented(
+ "TFLite quantized LSTM is not supported yet."
+ )
+
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 24, (
+ f"input tensors length should be 24, got {len(input_tensors)}"
+ )
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) >= 1, "output tensors length should be at
least 1"
+
+ assert op.BuiltinOptionsType() == BuiltinOptions.LSTMOptions
+ op_options = op.BuiltinOptions()
+ lstm_opts = LSTMOptions()
+ lstm_opts.Init(op_options.Bytes, op_options.Pos)
+
+ fused_activation_fn = lstm_opts.FusedActivationFunction()
+ cell_clip = lstm_opts.CellClip()
+ proj_clip = lstm_opts.ProjClip()
+
+ in_expr = self.get_tensor_expr(input_tensors[0])
+
+ # Only coupled input-forget gate is supported.
+ assert input_tensors[1].tensor_idx == -1, (
+ "Only coupled input-forget LSTM is supported."
+ )
+ assert input_tensors[5].tensor_idx == -1, (
+ "Only coupled input-forget LSTM is supported."
+ )
+
+ # Weights.
+ w_f = self.get_tensor_expr(input_tensors[2])
+ w_c = self.get_tensor_expr(input_tensors[3])
+ w_o = self.get_tensor_expr(input_tensors[4])
+
+ r_f = self.get_tensor_expr(input_tensors[6])
+ r_c = self.get_tensor_expr(input_tensors[7])
+ r_o = self.get_tensor_expr(input_tensors[8])
+
+ # Biases.
+ b_f = self.get_tensor_expr(input_tensors[13])
+ b_c = self.get_tensor_expr(input_tensors[14])
+ b_o = self.get_tensor_expr(input_tensors[15])
+
+ # State inputs.
+ h_prev = self.get_tensor_expr(input_tensors[18])
+ c_prev = self.get_tensor_expr(input_tensors[19])
+
+ # Coupled input-forget gate.
+ f = relax.op.sigmoid(
+ relax.op.add(
+ relax.op.add(
+ relax.op.matmul(in_expr, relax.op.permute_dims(w_f)),
+ relax.op.matmul(h_prev, relax.op.permute_dims(r_f)),
+ ),
+ b_f,
+ )
+ )
+ i = relax.op.subtract(
+ relax.const(np.array(1.0, dtype="float32"), "float32"),
+ f,
+ )
+
+ # Cell candidate.
+ g = relax.op.tanh(
+ relax.op.add(
+ relax.op.add(
+ relax.op.matmul(in_expr, relax.op.permute_dims(w_c)),
+ relax.op.matmul(h_prev, relax.op.permute_dims(r_c)),
+ ),
+ b_c,
+ )
+ )
+
+ # Output gate.
+ o = relax.op.sigmoid(
+ relax.op.add(
+ relax.op.add(
+ relax.op.matmul(in_expr, relax.op.permute_dims(w_o)),
+ relax.op.matmul(h_prev, relax.op.permute_dims(r_o)),
+ ),
+ b_o,
+ )
+ )
+
+ # Cell state update with optional clipping.
+ c_new = relax.op.add(
+ relax.op.multiply(f, c_prev),
+ relax.op.multiply(i, g),
+ )
+ if cell_clip > 0:
+ c_new = relax.op.clip(c_new, -cell_clip, cell_clip)
+
+ # Hidden state.
+ h_new = relax.op.multiply(o, relax.op.tanh(c_new))
+ if proj_clip > 0:
+ h_new = relax.op.clip(h_new, -proj_clip, proj_clip)
+
+ return self.convert_fused_activation_function(h_new,
fused_activation_fn)
+
+ def convert_svdf(self, op):
+ """Convert TFLite SVDF (single-step).
+
+ Structured-Vectorized Bidirectional Filter for keyword spotting.
+
+ Inputs (5 tensors):
+ [0] input [batch, input_size]
+ [1] feature_weights [num_filters, input_size]
+ [2] time_weights [num_filters, rank]
+ [3] bias [num_filters] (optional)
+ [4] state [batch, num_filters * rank] (variable)
+
+ Output:
+ [0] output [batch, num_filters]
+
+ Computation:
+ feat = x @ W_feat.T # feature projection
+ state_r = reshape(state, [B, F, rank]) # reshape
+ time = sum(state_r * time_weights, axis=-1) # time filtering
+ out = activation(feat + time + bias)
+ """
+ from tflite.BuiltinOptions import BuiltinOptions
+ from tflite.SVDFOptions import SVDFOptions
+
+ if self.is_quantized(op):
+ raise tvm.error.OpNotImplemented(
+ "TFLite quantized SVDF is not supported yet."
+ )
+
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 5, (
+ f"input tensors length should be 5, got {len(input_tensors)}"
+ )
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) >= 1, "output tensors length should be at
least 1"
+
+ assert op.BuiltinOptionsType() == BuiltinOptions.SVDFOptions
+ op_options = op.BuiltinOptions()
+ svdf_opts = SVDFOptions()
+ svdf_opts.Init(op_options.Bytes, op_options.Pos)
+
+ rank = svdf_opts.Rank()
+ fused_activation_fn = svdf_opts.FusedActivationFunction()
+
+ in_expr = self.get_tensor_expr(input_tensors[0])
+ feat_weights = self.get_tensor_expr(input_tensors[1])
+ time_weights = self.get_tensor_expr(input_tensors[2])
+
+ batch_size = to_int_list(self.get_tensor_shape(input_tensors[0]))[0]
+ num_filters = to_int_list(self.get_tensor_shape(input_tensors[1]))[0]
+
+ # Feature projection: [batch, input_size] @ [input_size, num_filters]
+ feat = relax.op.matmul(in_expr, relax.op.permute_dims(feat_weights))
+
+ # Time filtering: reshape state -> weight -> reduce
+ state_expr = self.get_tensor_expr(input_tensors[4])
+ state_3d = relax.op.reshape(state_expr, (batch_size, num_filters,
rank))
+
+ # time_weights: [num_filters, rank], broadcast to [1, num_filters,
rank]
+ tw_3d = relax.op.reshape(time_weights, (1, num_filters, rank))
+ time_weighted = relax.op.multiply(state_3d, tw_3d)
+ time_output = relax.op.sum(time_weighted, axis=-1, keepdims=False)
+
+ # Add feature + time
+ result = relax.op.add(feat, time_output)
+
+ # Add bias if present
+ if input_tensors[3].tensor_idx != -1:
+ bias_expr = self.get_tensor_expr(input_tensors[3])
+ result = relax.op.add(result, bias_expr)
+
+ return self.convert_fused_activation_function(result,
fused_activation_fn)
Review Comment:
The SVDF operator is stateful and updates its state tensor at each step by
shifting the history left and appending the new feature projection `feat`.
Without this state update, the operator is mathematically incorrect for any
sequence length greater than 1. We should compute the updated state and update
it in the `ExprTable` so that subsequent steps correctly receive the updated
state.
```python
# Add feature + time
result = relax.op.add(feat, time_output)
# Add bias if present
if input_tensors[3].tensor_idx != -1:
bias_expr = self.get_tensor_expr(input_tensors[3])
result = relax.op.add(result, bias_expr)
# Update state tensor in the expression table for subsequent ops
feat_3d = relax.op.expand_dims(feat, axis=-1)
if rank > 1:
shifted_state = relax.op.strided_slice(
state_3d, axes=[2], begin=[1], end=[rank]
)
new_state_3d = relax.op.concat([shifted_state, feat_3d], axis=2)
else:
new_state_3d = feat_3d
new_state = relax.op.reshape(new_state_3d, (batch_size, num_filters
* rank))
self.exp_tab.set_expr(
get_tensor_name(self.subgraph, input_tensors[4].tensor_idx),
new_state,
force_override=True,
)
return self.convert_fused_activation_function(result,
fused_activation_fn)
```
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -4878,6 +4880,222 @@ def convert_unpack(self, op):
return squeezed
+ def convert_lstm(self, op):
+ """Convert TFLite LSTM (single-step).
+
+ Standard LSTM cell with FULL kernel and coupled input-forget gate.
+ Peephole, projection, and layer norm are not supported.
+
+ Inputs (24 tensors, many optional):
+ [0] input [batch, input_size]
+ [1] input_to_input_weights (optional, -1 => coupled)
+ [2] input_to_forget_weights [num_units, input_size]
+ [3] input_to_cell_weights [num_units, input_size]
+ [4] input_to_output_weights [num_units, input_size]
+ [5] recurrent_to_input_weights (optional)
+ [6] recurrent_to_forget_weights [num_units, num_units]
+ [7] recurrent_to_cell_weights [num_units, num_units]
+ [8] recurrent_to_output_weights [num_units, num_units]
+ [9-11] cell_to_*_weights (optional, not supported)
+ [12] input_gate_bias (optional)
+ [13] forget_gate_bias [num_units]
+ [14] cell_bias [num_units]
+ [15] output_gate_bias [num_units]
+ [16-17] projection_weights/bias (optional, not supported)
+ [18] output_state [batch, num_units]
+ [19] cell_state [batch, num_units]
+ [20-23] layer_norm (optional, not supported)
+
+ Output:
+ [0] output [batch, num_units]
+
+ Cell (coupled input-forget):
+ f = sigmoid(x @ W_f.T + h @ R_f.T + b_f)
+ i = 1 - f
+ g = tanh(x @ W_c.T + h @ R_c.T + b_c)
+ o = sigmoid(x @ W_o.T + h @ R_o.T + b_o)
+ c_new = f * c_prev + i * g
+ h_new = fused_activation(o * tanh(c_new))
+ """
+ from tflite.BuiltinOptions import BuiltinOptions
+ from tflite.LSTMOptions import LSTMOptions
+
+ if self.is_quantized(op):
+ raise tvm.error.OpNotImplemented(
+ "TFLite quantized LSTM is not supported yet."
+ )
+
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 24, (
+ f"input tensors length should be 24, got {len(input_tensors)}"
+ )
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) >= 1, "output tensors length should be at
least 1"
+
+ assert op.BuiltinOptionsType() == BuiltinOptions.LSTMOptions
+ op_options = op.BuiltinOptions()
+ lstm_opts = LSTMOptions()
+ lstm_opts.Init(op_options.Bytes, op_options.Pos)
+
+ fused_activation_fn = lstm_opts.FusedActivationFunction()
+ cell_clip = lstm_opts.CellClip()
+ proj_clip = lstm_opts.ProjClip()
+
+ in_expr = self.get_tensor_expr(input_tensors[0])
+
+ # Only coupled input-forget gate is supported.
+ assert input_tensors[1].tensor_idx == -1, (
+ "Only coupled input-forget LSTM is supported."
+ )
+ assert input_tensors[5].tensor_idx == -1, (
+ "Only coupled input-forget LSTM is supported."
+ )
Review Comment:
Using `assert` statements to validate model properties or unsupported
features is discouraged because assertions can be disabled in Python when run
with optimization flags (e.g., `python -O`). This would bypass the check and
lead to silent failures or cryptic errors later. It is better to raise a proper
`tvm.error.OpNotImplemented` exception.
```python
# Only coupled input-forget gate is supported.
if input_tensors[1].tensor_idx != -1 or input_tensors[5].tensor_idx
!= -1:
raise tvm.error.OpNotImplemented(
"Only coupled input-forget LSTM is supported."
)
```
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
