Lines Matching refs:n_batch
305 const uint32_t n_batch = SizeOfDimension(input_, 0); in Prepare() local
333 outputShape->dimensions = {n_batch, n_output}; in Prepare()
338 outputStateShape->dimensions = {n_batch, n_output}; in Prepare()
343 cellStateShape->dimensions = {n_batch, n_cell}; in Prepare()
349 scratchShape->dimensions = {n_batch, n_cell * 3}; in Prepare()
352 scratchShape->dimensions = {n_batch, n_cell * 4}; in Prepare()
745 const uint32_t n_batch = input_shape.dimensions[0]; in LSTMStep() local
759 forget_gate_scratch = cell_scratch + n_cell * n_batch; in LSTMStep()
760 output_gate_scratch = cell_scratch + 2 * n_cell * n_batch; in LSTMStep()
763 cell_scratch = input_gate_scratch + n_cell * n_batch; in LSTMStep()
764 forget_gate_scratch = input_gate_scratch + 2 * n_cell * n_batch; in LSTMStep()
765 output_gate_scratch = input_gate_scratch + 3 * n_cell * n_batch; in LSTMStep()
771 tflite::tensor_utils::VectorBatchVectorAssign(input_gate_bias_buffer, n_cell, n_batch, in LSTMStep()
774 tflite::tensor_utils::VectorBatchVectorAssign(forget_gate_bias_buffer, n_cell, n_batch, in LSTMStep()
776 tflite::tensor_utils::VectorBatchVectorAssign(cell_bias_buffer, n_cell, n_batch, in LSTMStep()
778 tflite::tensor_utils::VectorBatchVectorAssign(output_gate_bias_buffer, n_cell, n_batch, in LSTMStep()
783 tflite::tensor_utils::ZeroVector(input_gate_scratch, n_cell * n_batch); in LSTMStep()
785 tflite::tensor_utils::ZeroVector(forget_gate_scratch, n_cell * n_batch); in LSTMStep()
786 tflite::tensor_utils::ZeroVector(cell_scratch, n_cell * n_batch); in LSTMStep()
787 tflite::tensor_utils::ZeroVector(output_gate_scratch, n_cell * n_batch); in LSTMStep()
793 input_to_input_weights_buffer, n_cell, n_input, input_buffer, n_batch, in LSTMStep()
797 input_to_forget_weights_buffer, n_cell, n_input, input_buffer, n_batch, in LSTMStep()
800 n_input, input_buffer, n_batch, in LSTMStep()
803 input_to_output_weights_buffer, n_cell, n_input, input_buffer, n_batch, in LSTMStep()
811 n_batch, input_gate_scratch, in LSTMStep()
816 aux_input_to_forget_weights_buffer, n_cell, n_aux_input, aux_input_buffer, n_batch, in LSTMStep()
819 aux_input_to_cell_weights_buffer, n_cell, n_aux_input, aux_input_buffer, n_batch, in LSTMStep()
822 aux_input_to_output_weights_buffer, n_cell, n_aux_input, aux_input_buffer, n_batch, in LSTMStep()
830 n_batch, input_gate_scratch, in LSTMStep()
834 recurrent_to_forget_weights_buffer, n_cell, n_output, output_state_in_buffer, n_batch, in LSTMStep()
837 recurrent_to_cell_weights_buffer, n_cell, n_output, output_state_in_buffer, n_batch, in LSTMStep()
840 recurrent_to_output_weights_buffer, n_cell, n_output, output_state_in_buffer, n_batch, in LSTMStep()
847 cell_to_input_weights_buffer, n_cell, cell_state_in_buffer, n_batch, in LSTMStep()
852 n_cell, n_batch, kLayerNormEpsilon); in LSTMStep()
854 n_cell, input_gate_scratch, n_batch, in LSTMStep()
856 tflite::tensor_utils::VectorBatchVectorAdd(input_gate_bias_buffer, n_cell, n_batch, in LSTMStep()
859 tflite::tensor_utils::ApplySigmoidToVector(input_gate_scratch, n_cell * n_batch, in LSTMStep()
867 n_batch, forget_gate_scratch); in LSTMStep()
871 n_cell, n_batch, kLayerNormEpsilon); in LSTMStep()
873 n_cell, forget_gate_scratch, n_batch, in LSTMStep()
875 tflite::tensor_utils::VectorBatchVectorAdd(forget_gate_bias_buffer, n_cell, n_batch, in LSTMStep()
878 tflite::tensor_utils::ApplySigmoidToVector(forget_gate_scratch, n_cell * n_batch, in LSTMStep()
883 tflite::tensor_utils::MeanStddevNormalization(cell_scratch, cell_scratch, n_cell, n_batch, in LSTMStep()
886 cell_scratch, n_batch, cell_scratch); in LSTMStep()
887 tflite::tensor_utils::VectorBatchVectorAdd(cell_bias_buffer, n_cell, n_batch, cell_scratch); in LSTMStep()
890 n_batch * n_cell, cell_state_out_buffer); in LSTMStep()
891 tflite::tensor_utils::ApplyActivationToVector(cell_scratch, n_batch * n_cell, params.activation, in LSTMStep()
894 tflite::tensor_utils::Sub1Vector(forget_gate_scratch, n_batch * n_cell, in LSTMStep()
897 cell_scratch, forget_gate_scratch, n_batch * n_cell, cell_state_out_buffer); in LSTMStep()
900 cell_scratch, input_gate_scratch, n_batch * n_cell, cell_state_out_buffer); in LSTMStep()
903 tflite::tensor_utils::ClipVector(cell_state_out_buffer, n_batch * n_cell, params.cell_clip, in LSTMStep()
911 n_batch, output_gate_scratch); in LSTMStep()
915 n_cell, n_batch, kLayerNormEpsilon); in LSTMStep()
917 n_cell, output_gate_scratch, n_batch, in LSTMStep()
919 tflite::tensor_utils::VectorBatchVectorAdd(output_gate_bias_buffer, n_cell, n_batch, in LSTMStep()
922 tflite::tensor_utils::ApplySigmoidToVector(output_gate_scratch, n_batch * n_cell, in LSTMStep()
924 tflite::tensor_utils::ApplyActivationToVector(cell_state_out_buffer, n_batch * n_cell, in LSTMStep()
927 n_batch * n_cell, output_gate_scratch); in LSTMStep()
932 tflite::tensor_utils::VectorBatchVectorAssign(projection_bias_buffer, n_output, n_batch, in LSTMStep()
935 tflite::tensor_utils::ZeroVector(output_buffer, n_batch * n_output); in LSTMStep()
938 projection_weights_buffer, n_output, n_cell, output_gate_scratch, n_batch, in LSTMStep()
942 tflite::tensor_utils::ClipVector(output_buffer, n_batch * n_output, params.proj_clip, in LSTMStep()
946 tflite::tensor_utils::CopyVector(output_gate_scratch, n_batch * n_output, output_buffer); in LSTMStep()
948 tflite::tensor_utils::CopyVector(output_buffer, n_batch * n_output, output_state_out_buffer); in LSTMStep()