Method Initialize

tensorflow/compiler/mlir/lite/utils/lstm_utils.cc:438–482

}

LogicalResult ConvertLSTMCellSimpleToFusedLSTM::Initialize() {
  if (failed(InitializeFromFuncAttributes())) {
    return fused_func_op_.emitError()
           << "Expected function attributes were not set on the function "
              "encapsulating the composite op";
  }

  num_gates_ = couple_input_forget_gates_ ? 3 : 4;

  input_ = fused_func_op_.getArgument(0);
  bias_ = fused_func_op_.getArgument(2);

  weight_ = fused_func_op_.getArgument(1);
  weight_type_ = weight_.getType().cast<RankedTensorType>();

  if (weight_type_.getRank() != 2) {
    return fused_func_op_.emitError() << "The weight tensor was not of rank 2";
  }

  if (weight_type_.getDimSize(1) % num_gates_ != 0) {
    return fused_func_op_.emitError()
           << "Invalid dimension 1 of weight tensor, "
              "should be divisible by the number of gates";
  }
  n_cell_ = weight_type_.getDimSize(1) / num_gates_;

  projection_ = fused_func_op_.getArgument(3);
  projection_type_ = projection_.getType().cast<RankedTensorType>();
  if (projection_type_.getRank() != 2) {
    n_output_ = n_cell_;
  } else {
    n_output_ = projection_type_.getDimSize(1);
  }
  n_input_ = weight_type_.getDimSize(0) - n_output_;
  num_cols_weight_transposed_ = weight_type_.getDimSize(0);
  num_cols_projection_transposed_ = projection_type_.getDimSize(0);

  bias_slice_shape_ = {n_cell_};
  bias_size_values_ = {n_cell_};
  weight_slice_shape_ = {1, num_cols_weight_transposed_};
  weight_slice_size_input_values_ = {n_cell_, n_input_};
  weight_slice_size_recurrent_values_ = {n_cell_, n_output_};

  return success();
}

LogicalResult ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM::Initialize() {
  if (failed(ConvertLSTMCellSimpleToFusedLSTM::Initialize())) {