leduclinh's picture
Duplicate from aufklarer/Parakeet-TDT-v3-CoreML-INT4
e548cac
program(1.0)
[buildInfo = dict<tensor<string, []>, tensor<string, []>>({{"coremlc-component-MIL", "3510.2.1"}, {"coremlc-version", "3500.32.1"}, {"coremltools-component-torch", "2.10.0"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "8.1"}})]
{
func main<ios17>(tensor<fp16, [2, 1, 640]> c, tensor<fp16, [2, 1, 640]> h, tensor<int32, [1, 1]> token) {
tensor<int32, []> y_1_axis_0 = const()[name = tensor<string, []>("y_1_axis_0"), val = tensor<int32, []>(0)];
tensor<int32, []> y_1_batch_dims_0 = const()[name = tensor<string, []>("y_1_batch_dims_0"), val = tensor<int32, []>(0)];
tensor<bool, []> y_1_validate_indices_0 = const()[name = tensor<string, []>("y_1_validate_indices_0"), val = tensor<bool, []>(false)];
tensor<fp16, [8193, 640]> decoder_prediction_embed_weight_to_fp16 = const()[name = tensor<string, []>("decoder_prediction_embed_weight_to_fp16"), val = tensor<fp16, [8193, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(64)))];
tensor<string, []> token_to_int16_dtype_0 = const()[name = tensor<string, []>("token_to_int16_dtype_0"), val = tensor<string, []>("int16")];
tensor<int16, [1, 1]> token_to_int16 = cast(dtype = token_to_int16_dtype_0, x = token)[name = tensor<string, []>("cast_6")];
tensor<fp16, [1, 1, 640]> y_1_cast_fp16_cast_uint16 = gather(axis = y_1_axis_0, batch_dims = y_1_batch_dims_0, indices = token_to_int16, validate_indices = y_1_validate_indices_0, x = decoder_prediction_embed_weight_to_fp16)[name = tensor<string, []>("y_1_cast_fp16_cast_uint16")];
tensor<int32, [3]> input_1_perm_0 = const()[name = tensor<string, []>("input_1_perm_0"), val = tensor<int32, [3]>([1, 0, 2])];
tensor<int32, []> split_0_num_splits_0 = const()[name = tensor<string, []>("split_0_num_splits_0"), val = tensor<int32, []>(2)];
tensor<int32, []> split_0_axis_0 = const()[name = tensor<string, []>("split_0_axis_0"), val = tensor<int32, []>(0)];
tensor<fp16, [1, 1, 640]> split_0_cast_fp16_0, tensor<fp16, [1, 1, 640]> split_0_cast_fp16_1 = split(axis = split_0_axis_0, num_splits = split_0_num_splits_0, x = h)[name = tensor<string, []>("split_0_cast_fp16")];
tensor<int32, []> split_1_num_splits_0 = const()[name = tensor<string, []>("split_1_num_splits_0"), val = tensor<int32, []>(2)];
tensor<int32, []> split_1_axis_0 = const()[name = tensor<string, []>("split_1_axis_0"), val = tensor<int32, []>(0)];
tensor<fp16, [1, 1, 640]> split_1_cast_fp16_0, tensor<fp16, [1, 1, 640]> split_1_cast_fp16_1 = split(axis = split_1_axis_0, num_splits = split_1_num_splits_0, x = c)[name = tensor<string, []>("split_1_cast_fp16")];
tensor<int32, [1]> input0_1_lstm_layer_0_lstm_h0_squeeze_axes_0 = const()[name = tensor<string, []>("input0_1_lstm_layer_0_lstm_h0_squeeze_axes_0"), val = tensor<int32, [1]>([0])];
tensor<fp16, [1, 640]> input0_1_lstm_layer_0_lstm_h0_squeeze_cast_fp16 = squeeze(axes = input0_1_lstm_layer_0_lstm_h0_squeeze_axes_0, x = split_0_cast_fp16_0)[name = tensor<string, []>("input0_1_lstm_layer_0_lstm_h0_squeeze_cast_fp16")];
tensor<int32, [1]> input0_1_lstm_layer_0_lstm_c0_squeeze_axes_0 = const()[name = tensor<string, []>("input0_1_lstm_layer_0_lstm_c0_squeeze_axes_0"), val = tensor<int32, [1]>([0])];
tensor<fp16, [1, 640]> input0_1_lstm_layer_0_lstm_c0_squeeze_cast_fp16 = squeeze(axes = input0_1_lstm_layer_0_lstm_c0_squeeze_axes_0, x = split_1_cast_fp16_0)[name = tensor<string, []>("input0_1_lstm_layer_0_lstm_c0_squeeze_cast_fp16")];
tensor<string, []> input0_1_lstm_layer_0_direction_0 = const()[name = tensor<string, []>("input0_1_lstm_layer_0_direction_0"), val = tensor<string, []>("forward")];
tensor<bool, []> input0_1_lstm_layer_0_output_sequence_0 = const()[name = tensor<string, []>("input0_1_lstm_layer_0_output_sequence_0"), val = tensor<bool, []>(true)];
tensor<string, []> input0_1_lstm_layer_0_recurrent_activation_0 = const()[name = tensor<string, []>("input0_1_lstm_layer_0_recurrent_activation_0"), val = tensor<string, []>("sigmoid")];
tensor<string, []> input0_1_lstm_layer_0_cell_activation_0 = const()[name = tensor<string, []>("input0_1_lstm_layer_0_cell_activation_0"), val = tensor<string, []>("tanh")];
tensor<string, []> input0_1_lstm_layer_0_activation_0 = const()[name = tensor<string, []>("input0_1_lstm_layer_0_activation_0"), val = tensor<string, []>("tanh")];
tensor<fp16, [2560, 640]> concat_1_to_fp16 = const()[name = tensor<string, []>("concat_1_to_fp16"), val = tensor<fp16, [2560, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(10487168)))];
tensor<fp16, [2560, 640]> concat_2_to_fp16 = const()[name = tensor<string, []>("concat_2_to_fp16"), val = tensor<fp16, [2560, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(13764032)))];
tensor<fp16, [2560]> concat_0_to_fp16 = const()[name = tensor<string, []>("concat_0_to_fp16"), val = tensor<fp16, [2560]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(17040896)))];
tensor<fp16, [1, 1, 640]> input_1_cast_fp16 = transpose(perm = input_1_perm_0, x = y_1_cast_fp16_cast_uint16)[name = tensor<string, []>("transpose_1")];
tensor<fp16, [1, 1, 640]> input0_1_lstm_layer_0_cast_fp16_0, tensor<fp16, [1, 640]> input0_1_lstm_layer_0_cast_fp16_1, tensor<fp16, [1, 640]> input0_1_lstm_layer_0_cast_fp16_2 = lstm(activation = input0_1_lstm_layer_0_activation_0, bias = concat_0_to_fp16, cell_activation = input0_1_lstm_layer_0_cell_activation_0, direction = input0_1_lstm_layer_0_direction_0, initial_c = input0_1_lstm_layer_0_lstm_c0_squeeze_cast_fp16, initial_h = input0_1_lstm_layer_0_lstm_h0_squeeze_cast_fp16, output_sequence = input0_1_lstm_layer_0_output_sequence_0, recurrent_activation = input0_1_lstm_layer_0_recurrent_activation_0, weight_hh = concat_2_to_fp16, weight_ih = concat_1_to_fp16, x = input_1_cast_fp16)[name = tensor<string, []>("input0_1_lstm_layer_0_cast_fp16")];
tensor<int32, [1]> input0_1_lstm_h0_squeeze_axes_0 = const()[name = tensor<string, []>("input0_1_lstm_h0_squeeze_axes_0"), val = tensor<int32, [1]>([0])];
tensor<fp16, [1, 640]> input0_1_lstm_h0_squeeze_cast_fp16 = squeeze(axes = input0_1_lstm_h0_squeeze_axes_0, x = split_0_cast_fp16_1)[name = tensor<string, []>("input0_1_lstm_h0_squeeze_cast_fp16")];
tensor<int32, [1]> input0_1_lstm_c0_squeeze_axes_0 = const()[name = tensor<string, []>("input0_1_lstm_c0_squeeze_axes_0"), val = tensor<int32, [1]>([0])];
tensor<fp16, [1, 640]> input0_1_lstm_c0_squeeze_cast_fp16 = squeeze(axes = input0_1_lstm_c0_squeeze_axes_0, x = split_1_cast_fp16_1)[name = tensor<string, []>("input0_1_lstm_c0_squeeze_cast_fp16")];
tensor<string, []> input0_1_direction_0 = const()[name = tensor<string, []>("input0_1_direction_0"), val = tensor<string, []>("forward")];
tensor<bool, []> input0_1_output_sequence_0 = const()[name = tensor<string, []>("input0_1_output_sequence_0"), val = tensor<bool, []>(true)];
tensor<string, []> input0_1_recurrent_activation_0 = const()[name = tensor<string, []>("input0_1_recurrent_activation_0"), val = tensor<string, []>("sigmoid")];
tensor<string, []> input0_1_cell_activation_0 = const()[name = tensor<string, []>("input0_1_cell_activation_0"), val = tensor<string, []>("tanh")];
tensor<string, []> input0_1_activation_0 = const()[name = tensor<string, []>("input0_1_activation_0"), val = tensor<string, []>("tanh")];
tensor<fp16, [2560, 640]> concat_4_to_fp16 = const()[name = tensor<string, []>("concat_4_to_fp16"), val = tensor<fp16, [2560, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(17046080)))];
tensor<fp16, [2560, 640]> concat_5_to_fp16 = const()[name = tensor<string, []>("concat_5_to_fp16"), val = tensor<fp16, [2560, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(20322944)))];
tensor<fp16, [2560]> concat_3_to_fp16 = const()[name = tensor<string, []>("concat_3_to_fp16"), val = tensor<fp16, [2560]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(23599808)))];
tensor<fp16, [1, 1, 640]> input0_1_cast_fp16_0, tensor<fp16, [1, 640]> input0_1_cast_fp16_1, tensor<fp16, [1, 640]> input0_1_cast_fp16_2 = lstm(activation = input0_1_activation_0, bias = concat_3_to_fp16, cell_activation = input0_1_cell_activation_0, direction = input0_1_direction_0, initial_c = input0_1_lstm_c0_squeeze_cast_fp16, initial_h = input0_1_lstm_h0_squeeze_cast_fp16, output_sequence = input0_1_output_sequence_0, recurrent_activation = input0_1_recurrent_activation_0, weight_hh = concat_5_to_fp16, weight_ih = concat_4_to_fp16, x = input0_1_lstm_layer_0_cast_fp16_0)[name = tensor<string, []>("input0_1_cast_fp16")];
tensor<int32, []> var_33_axis_0 = const()[name = tensor<string, []>("op_33_axis_0"), val = tensor<int32, []>(0)];
tensor<fp16, [2, 1, 640]> h_out = stack(axis = var_33_axis_0, values = (input0_1_lstm_layer_0_cast_fp16_1, input0_1_cast_fp16_1))[name = tensor<string, []>("op_33_cast_fp16")];
tensor<int32, []> var_34_axis_0 = const()[name = tensor<string, []>("op_34_axis_0"), val = tensor<int32, []>(0)];
tensor<fp16, [2, 1, 640]> c_out = stack(axis = var_34_axis_0, values = (input0_1_lstm_layer_0_cast_fp16_2, input0_1_cast_fp16_2))[name = tensor<string, []>("op_34_cast_fp16")];
tensor<int32, [3]> var_44_perm_0 = const()[name = tensor<string, []>("op_44_perm_0"), val = tensor<int32, [3]>([1, 0, 2])];
tensor<fp16, [1, 1, 640]> decoder_output = transpose(perm = var_44_perm_0, x = input0_1_cast_fp16_0)[name = tensor<string, []>("transpose_0")];
} -> (decoder_output, h_out, c_out);
}