bidirectional_sequence_lstm.cc | Explore in Territory

/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <math.h>

#include <algorithm>
#include <cstddef>

#include "tensorflow/lite/core/c/builtin_op_data.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/internal/kernel_utils.h"
#include "tensorflow/lite/kernels/internal/tensor_utils.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/lstm_eval.h"
#include "tensorflow/lite/kernels/op_macros.h"

namespace tflite {
namespace ops {
namespace builtin {
namespace bidirectional_sequence_lstm {

// LINT.IfChange

// Input Tensors of size {max_time, n_batch, n_input}
constexpr int kInputTensor = …;

// Forward LSTM cell tensors.
// Input weight tensors of size: {n_cell, n_input}
constexpr int kFwInputToInputWeightsTensor = …;  // Optional
constexpr int kFwInputToForgetWeightsTensor = …;
constexpr int kFwInputToCellWeightsTensor = …;
constexpr int kFwInputToOutputWeightsTensor = …;

// Recurrent weight tensors of size {n_cell, n_output}
constexpr int kFwRecurrentToInputWeightsTensor = …;  // Optional
constexpr int kFwRecurrentToForgetWeightsTensor = …;
constexpr int kFwRecurrentToCellWeightsTensor = …;
constexpr int kFwRecurrentToOutputWeightsTensor = …;

// Peephole weights tensors of size {n_cell}, representing a diagonal matrix.
constexpr int kFwCellToInputWeightsTensor = …;    // Optional
constexpr int kFwCellToForgetWeightsTensor = …;  // Optional
constexpr int kFwCellToOutputWeightsTensor = …;  // Optional

// Gates bias tensors of size {n_cell}
constexpr int kFwInputGateBiasTensor = …;  // Optional
constexpr int kFwForgetGateBiasTensor = …;
constexpr int kFwCellGateBiasTensor = …;
constexpr int kFwOutputGateBiasTensor = …;

// Projection weight tensor of size {n_output, n_cell}
constexpr int kFwProjectionWeightsTensor = …;  // Optional
// Projection bias tensor of size {n_output}
constexpr int kFwProjectionBiasTensor = …;  // Optional

// Backward LSTM cell tensors.
// Input weight tensors of size: {n_cell, n_input}
constexpr int kBwInputToInputWeightsTensor = …;  // Optional
constexpr int kBwInputToForgetWeightsTensor = …;
constexpr int kBwInputToCellWeightsTensor = …;
constexpr int kBwInputToOutputWeightsTensor = …;

// Recurrent weight tensors of size {n_cell, n_output}
constexpr int kBwRecurrentToInputWeightsTensor = …;  // Optional
constexpr int kBwRecurrentToForgetWeightsTensor = …;
constexpr int kBwRecurrentToCellWeightsTensor = …;
constexpr int kBwRecurrentToOutputWeightsTensor = …;

// Peephole weights tensors of size {n_cell}, representing a diagonal matrix.
constexpr int kBwCellToInputWeightsTensor = …;   // Optional
constexpr int kBwCellToForgetWeightsTensor = …;  // Optional
constexpr int kBwCellToOutputWeightsTensor = …;  // Optional

// Gates bias tensors of size {n_cell}
constexpr int kBwInputGateBiasTensor = …;  // Optional
constexpr int kBwForgetGateBiasTensor = …;
constexpr int kBwCellGateBiasTensor = …;
constexpr int kBwOutputGateBiasTensor = …;

// Projection weight tensor of size {n_output, n_cell}
constexpr int kBwProjectionWeightsTensor = …;  // Optional
// Projection bias tensor of size {n_output}
constexpr int kBwProjectionBiasTensor = …;  // Optional

// Stateful input tensors that are variables and will be modified by the Op.
// Activation state tensors of size {n_batch, n_output}
constexpr int kFwInputActivationStateTensor = …;
// Cell state tensors of size {n_batch, n_cell}
constexpr int kFwInputCellStateTensor = …;
// Activation state tensors of size {n_batch, n_output}
constexpr int kBwInputActivationStateTensor = …;
// Cell state tensors of size {n_batch, n_cell}
constexpr int kBwInputCellStateTensor = …;

// Used as auxiliary input and weights when stacking for
// tf.contrib.rnn.stack_bidirectional_rnn case (with cross links); Used as input
// to the backward cell when stacking for tf.nn.static_bidirectional_rnn case
// (without cross links).
constexpr int kAuxInputTensor = …;  // Optional
// Forward weights.
constexpr int kFwAuxInputToInputWeightsTensor = …;   // Optional
constexpr int kFwAuxInputToForgetWeightsTensor = …;  // Optional
constexpr int kFwAuxInputToCellWeightsTensor = …;    // Optional
constexpr int kFwAuxInputToOutputWeightsTensor = …;  // Optional
// Backward weights.
constexpr int kBwAuxInputToInputWeightsTensor = …;   // Optional
constexpr int kBwAuxInputToForgetWeightsTensor = …;  // Optional
constexpr int kBwAuxInputToCellWeightsTensor = …;    // Optional
constexpr int kBwAuxInputToOutputWeightsTensor = …;  // Optional

// Output tensors.
constexpr int kFwOutputTensor = …;
constexpr int kBwOutputTensor = …;  // Ignored if merge_outputs is set.

// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantize_weights.cc)

// Temporary tensors.
enum TemporaryTensor { … };

struct OpData { … };

void* Init(TfLiteContext* context, const char* buffer, size_t length) { … }

void Free(TfLiteContext* context, void* buffer) { … }

// Check that input tensor dimensions matches with each other.
TfLiteStatus CheckLstmTensorDimensionsAndTypes(
    TfLiteContext* context, TfLiteNode* node, int n_input, int n_output,
    int n_cell, int input_to_input_weights_tensor,
    int input_to_forget_weights_tensor, int input_to_cell_weights_tensor,
    int input_to_output_weights_tensor, int recurrent_to_input_weights_tensor,
    int recurrent_to_forget_weights_tensor,
    int recurrent_to_cell_weights_tensor,
    int recurrent_to_output_weights_tensor, int cell_to_input_weights_tensor,
    int cell_to_forget_weights_tensor, int cell_to_output_weights_tensor,
    int input_gate_bias_tensor, int forget_gate_bias_tensor,
    int cell_gate_bias_tensor, int output_gate_bias_tensor,
    int projection_weights_tensor, int projection_bias_tensor) { … }

TfLiteStatus CheckInputTensorDimensions(TfLiteContext* context,
                                        TfLiteNode* node, int n_input,
                                        int n_output, int n_cell) { … }

// Resize the output and scratch tensors based on the sizes of the input
// tensors. Also check that the size of the input tensors match each other.
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { … }

// The LSTM Op engine.
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { … }

}  // namespace bidirectional_sequence_lstm

TfLiteRegistration* Register_BIDIRECTIONAL_SEQUENCE_LSTM() { … }

}  // namespace builtin
}  // namespace ops
}  // namespace tflite
chromium/third_party/tflite/src/tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
