| program(1.0) |
| [buildInfo = dict<tensor<string, []>, tensor<string, []>>({{"coremlc-component-MIL", "3405.2.1"}, {"coremlc-version", "3404.23.1"}, {"coremltools-component-torch", "2.7.0"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "9.0b1"}})] |
| { |
| func main<ios17>(tensor<fp32, [1, 640, 1]> decoder_step, tensor<fp32, [1, 1024, 1]> encoder_step) { |
| tensor<fp32, [640]> joint_module_enc_bias = const()[name = tensor<string, []>("joint_module_enc_bias"), val = tensor<fp32, [640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(64)))]; |
| tensor<fp32, [640, 1024]> joint_module_enc_weight = const()[name = tensor<string, []>("joint_module_enc_weight"), val = tensor<fp32, [640, 1024]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(2688)))]; |
| tensor<fp32, [640]> joint_module_pred_bias = const()[name = tensor<string, []>("joint_module_pred_bias"), val = tensor<fp32, [640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(2624192)))]; |
| tensor<fp32, [640, 640]> joint_module_pred_weight = const()[name = tensor<string, []>("joint_module_pred_weight"), val = tensor<fp32, [640, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(2626816)))]; |
| tensor<fp32, [1030]> joint_module_joint_net_2_bias = const()[name = tensor<string, []>("joint_module_joint_net_2_bias"), val = tensor<fp32, [1030]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(4265280)))]; |
| tensor<fp32, [1030, 640]> joint_module_joint_net_2_weight = const()[name = tensor<string, []>("joint_module_joint_net_2_weight"), val = tensor<fp32, [1030, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(4269504)))]; |
| tensor<int32, [3]> input_1_perm_0 = const()[name = tensor<string, []>("input_1_perm_0"), val = tensor<int32, [3]>([0, 2, 1])]; |
| tensor<int32, [3]> input_3_perm_0 = const()[name = tensor<string, []>("input_3_perm_0"), val = tensor<int32, [3]>([0, 2, 1])]; |
| tensor<fp32, [1, 1, 1024]> input_1 = transpose(perm = input_1_perm_0, x = encoder_step)[name = tensor<string, []>("transpose_1")]; |
| tensor<fp32, [1, 1, 640]> enc_proj = linear(bias = joint_module_enc_bias, weight = joint_module_enc_weight, x = input_1)[name = tensor<string, []>("linear_0")]; |
| tensor<fp32, [1, 1, 640]> input_3 = transpose(perm = input_3_perm_0, x = decoder_step)[name = tensor<string, []>("transpose_0")]; |
| tensor<fp32, [1, 1, 640]> dec_proj = linear(bias = joint_module_pred_bias, weight = joint_module_pred_weight, x = input_3)[name = tensor<string, []>("linear_1")]; |
| tensor<int32, [1]> var_23_axes_0 = const()[name = tensor<string, []>("op_23_axes_0"), val = tensor<int32, [1]>([2])]; |
| tensor<fp32, [1, 1, 1, 640]> var_23 = expand_dims(axes = var_23_axes_0, x = enc_proj)[name = tensor<string, []>("op_23")]; |
| tensor<int32, [1]> var_24_axes_0 = const()[name = tensor<string, []>("op_24_axes_0"), val = tensor<int32, [1]>([1])]; |
| tensor<fp32, [1, 1, 1, 640]> var_24 = expand_dims(axes = var_24_axes_0, x = dec_proj)[name = tensor<string, []>("op_24")]; |
| tensor<fp32, [1, 1, 1, 640]> input_5 = add(x = var_23, y = var_24)[name = tensor<string, []>("input_5")]; |
| tensor<fp32, [1, 1, 1, 640]> input_7 = relu(x = input_5)[name = tensor<string, []>("input_7")]; |
| tensor<fp32, [1, 1, 1, 1030]> logits = linear(bias = joint_module_joint_net_2_bias, weight = joint_module_joint_net_2_weight, x = input_7)[name = tensor<string, []>("linear_2")]; |
| tensor<int32, [4]> token_logits_begin_0 = const()[name = tensor<string, []>("token_logits_begin_0"), val = tensor<int32, [4]>([0, 0, 0, 0])]; |
| tensor<int32, [4]> token_logits_end_0 = const()[name = tensor<string, []>("token_logits_end_0"), val = tensor<int32, [4]>([1, 1, 1, 1025])]; |
| tensor<bool, [4]> token_logits_end_mask_0 = const()[name = tensor<string, []>("token_logits_end_mask_0"), val = tensor<bool, [4]>([true, true, true, false])]; |
| tensor<fp32, [1, 1, 1, 1025]> token_logits = slice_by_index(begin = token_logits_begin_0, end = token_logits_end_0, end_mask = token_logits_end_mask_0, x = logits)[name = tensor<string, []>("token_logits")]; |
| tensor<int32, [4]> duration_logits_begin_0 = const()[name = tensor<string, []>("duration_logits_begin_0"), val = tensor<int32, [4]>([0, 0, 0, 1025])]; |
| tensor<int32, [4]> duration_logits_end_0 = const()[name = tensor<string, []>("duration_logits_end_0"), val = tensor<int32, [4]>([1, 1, 1, 1030])]; |
| tensor<bool, [4]> duration_logits_end_mask_0 = const()[name = tensor<string, []>("duration_logits_end_mask_0"), val = tensor<bool, [4]>([true, true, true, true])]; |
| tensor<fp32, [1, 1, 1, 5]> duration_logits = slice_by_index(begin = duration_logits_begin_0, end = duration_logits_end_0, end_mask = duration_logits_end_mask_0, x = logits)[name = tensor<string, []>("duration_logits")]; |
| tensor<int32, []> var_43_axis_0 = const()[name = tensor<string, []>("op_43_axis_0"), val = tensor<int32, []>(-1)]; |
| tensor<bool, []> var_43_keep_dims_0 = const()[name = tensor<string, []>("op_43_keep_dims_0"), val = tensor<bool, []>(false)]; |
| tensor<string, []> var_43_output_dtype_0 = const()[name = tensor<string, []>("op_43_output_dtype_0"), val = tensor<string, []>("int32")]; |
| tensor<int32, [1, 1, 1]> token_id = reduce_argmax(axis = var_43_axis_0, keep_dims = var_43_keep_dims_0, output_dtype = var_43_output_dtype_0, x = token_logits)[name = tensor<string, []>("op_43")]; |
| tensor<int32, []> var_49 = const()[name = tensor<string, []>("op_49"), val = tensor<int32, []>(-1)]; |
| tensor<fp32, [1, 1, 1, 1025]> token_probs_all = softmax(axis = var_49, x = token_logits)[name = tensor<string, []>("token_probs_all")]; |
| tensor<int32, [1]> var_58_axes_0 = const()[name = tensor<string, []>("op_58_axes_0"), val = tensor<int32, [1]>([-1])]; |
| tensor<int32, [1, 1, 1, 1]> var_58 = expand_dims(axes = var_58_axes_0, x = token_id)[name = tensor<string, []>("op_58")]; |
| tensor<int32, []> var_59 = const()[name = tensor<string, []>("op_59"), val = tensor<int32, []>(-1)]; |
| tensor<bool, []> var_61_validate_indices_0 = const()[name = tensor<string, []>("op_61_validate_indices_0"), val = tensor<bool, []>(false)]; |
| tensor<fp32, [1, 1, 1, 1]> var_61 = gather_along_axis(axis = var_59, indices = var_58, validate_indices = var_61_validate_indices_0, x = token_probs_all)[name = tensor<string, []>("op_61")]; |
| tensor<int32, [1]> var_63_axes_0 = const()[name = tensor<string, []>("op_63_axes_0"), val = tensor<int32, [1]>([-1])]; |
| tensor<fp32, [1, 1, 1]> token_prob = squeeze(axes = var_63_axes_0, x = var_61)[name = tensor<string, []>("op_63")]; |
| tensor<int32, []> var_66_axis_0 = const()[name = tensor<string, []>("op_66_axis_0"), val = tensor<int32, []>(-1)]; |
| tensor<bool, []> var_66_keep_dims_0 = const()[name = tensor<string, []>("op_66_keep_dims_0"), val = tensor<bool, []>(false)]; |
| tensor<string, []> var_66_output_dtype_0 = const()[name = tensor<string, []>("op_66_output_dtype_0"), val = tensor<string, []>("int32")]; |
| tensor<int32, [1, 1, 1]> duration = reduce_argmax(axis = var_66_axis_0, keep_dims = var_66_keep_dims_0, output_dtype = var_66_output_dtype_0, x = duration_logits)[name = tensor<string, []>("op_66")]; |
| } -> (token_id, token_prob, duration); |
| } |
