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, [1, 1, 640]> decoder_output, tensor<fp16, [1, 1, 1024]> encoder_output) {
tensor<int32, []> var_6 = const()[name = tensor<string, []>("op_6"), val = tensor<int32, []>(-1)];
tensor<fp16, [640, 1024]> joint_enc_weight_to_fp16 = const()[name = tensor<string, []>("joint_enc_weight_to_fp16"), val = tensor<fp16, [640, 1024]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(64)))];
tensor<fp16, [640]> joint_enc_bias_to_fp16 = const()[name = tensor<string, []>("joint_enc_bias_to_fp16"), val = tensor<fp16, [640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(1310848)))];
tensor<fp16, [1, 1, 640]> linear_0_cast_fp16 = linear(bias = joint_enc_bias_to_fp16, weight = joint_enc_weight_to_fp16, x = encoder_output)[name = tensor<string, []>("linear_0_cast_fp16")];
tensor<fp16, [640, 640]> joint_pred_weight_to_fp16 = const()[name = tensor<string, []>("joint_pred_weight_to_fp16"), val = tensor<fp16, [640, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(1312192)))];
tensor<fp16, [640]> joint_pred_bias_to_fp16 = const()[name = tensor<string, []>("joint_pred_bias_to_fp16"), val = tensor<fp16, [640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(2131456)))];
tensor<fp16, [1, 1, 640]> linear_1_cast_fp16 = linear(bias = joint_pred_bias_to_fp16, weight = joint_pred_weight_to_fp16, x = decoder_output)[name = tensor<string, []>("linear_1_cast_fp16")];
tensor<int32, [1]> f_3_axes_0 = const()[name = tensor<string, []>("f_3_axes_0"), val = tensor<int32, [1]>([2])];
tensor<fp16, [1, 1, 1, 640]> f_3_cast_fp16 = expand_dims(axes = f_3_axes_0, x = linear_0_cast_fp16)[name = tensor<string, []>("f_3_cast_fp16")];
tensor<int32, [1]> g_3_axes_0 = const()[name = tensor<string, []>("g_3_axes_0"), val = tensor<int32, [1]>([1])];
tensor<fp16, [1, 1, 1, 640]> g_3_cast_fp16 = expand_dims(axes = g_3_axes_0, x = linear_1_cast_fp16)[name = tensor<string, []>("g_3_cast_fp16")];
tensor<fp16, [1, 1, 1, 640]> input_3_cast_fp16 = add(x = f_3_cast_fp16, y = g_3_cast_fp16)[name = tensor<string, []>("input_3_cast_fp16")];
tensor<fp16, [1, 1, 1, 640]> var_28_cast_fp16 = relu(x = input_3_cast_fp16)[name = tensor<string, []>("op_28_cast_fp16")];
tensor<fp16, [8198, 640]> joint_joint_net_2_weight_to_fp16 = const()[name = tensor<string, []>("joint_joint_net_2_weight_to_fp16"), val = tensor<fp16, [8198, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(2132800)))];
tensor<fp16, [8198]> joint_joint_net_2_bias_to_fp16 = const()[name = tensor<string, []>("joint_joint_net_2_bias_to_fp16"), val = tensor<fp16, [8198]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(12626304)))];
tensor<fp16, [1, 1, 1, 8198]> linear_2_cast_fp16 = linear(bias = joint_joint_net_2_bias_to_fp16, weight = joint_joint_net_2_weight_to_fp16, x = var_28_cast_fp16)[name = tensor<string, []>("linear_2_cast_fp16")];
tensor<fp16, [1, 1, 1, 8198]> combined_1_softmax_cast_fp16 = softmax(axis = var_6, x = linear_2_cast_fp16)[name = tensor<string, []>("combined_1_softmax_cast_fp16")];
tensor<fp32, []> combined_1_epsilon_0 = const()[name = tensor<string, []>("combined_1_epsilon_0"), val = tensor<fp32, []>(0x1p-149)];
tensor<fp16, [1, 1, 1, 8198]> combined_1_cast_fp16 = log(epsilon = combined_1_epsilon_0, x = combined_1_softmax_cast_fp16)[name = tensor<string, []>("combined_1_cast_fp16")];
tensor<int32, [1]> combined0_1_axes_0 = const()[name = tensor<string, []>("combined0_1_axes_0"), val = tensor<int32, [1]>([2])];
tensor<fp16, [1, 1, 8198]> combined0_1_cast_fp16 = squeeze(axes = combined0_1_axes_0, x = combined_1_cast_fp16)[name = tensor<string, []>("combined0_1_cast_fp16")];
tensor<int32, [3]> var_35_begin_0 = const()[name = tensor<string, []>("op_35_begin_0"), val = tensor<int32, [3]>([0, 0, 0])];
tensor<int32, [3]> var_35_end_0 = const()[name = tensor<string, []>("op_35_end_0"), val = tensor<int32, [3]>([1, 1, 8193])];
tensor<bool, [3]> var_35_end_mask_0 = const()[name = tensor<string, []>("op_35_end_mask_0"), val = tensor<bool, [3]>([true, true, false])];
tensor<fp16, [1, 1, 8193]> token_logits = slice_by_index(begin = var_35_begin_0, end = var_35_end_0, end_mask = var_35_end_mask_0, x = combined0_1_cast_fp16)[name = tensor<string, []>("op_35_cast_fp16")];
tensor<int32, [3]> var_36_begin_0 = const()[name = tensor<string, []>("op_36_begin_0"), val = tensor<int32, [3]>([0, 0, 8193])];
tensor<int32, [3]> var_36_end_0 = const()[name = tensor<string, []>("op_36_end_0"), val = tensor<int32, [3]>([1, 1, 8198])];
tensor<bool, [3]> var_36_end_mask_0 = const()[name = tensor<string, []>("op_36_end_mask_0"), val = tensor<bool, [3]>([true, true, true])];
tensor<fp16, [1, 1, 5]> duration_logits = slice_by_index(begin = var_36_begin_0, end = var_36_end_0, end_mask = var_36_end_mask_0, x = combined0_1_cast_fp16)[name = tensor<string, []>("op_36_cast_fp16")];
} -> (token_logits, duration_logits);
}