TokenDurationPrediction.mlmodelc/model.mil

program(1.0)
[buildInfo = dict<tensor<string, []>, tensor<string, []>>({{"coremlc-component-MIL", "3405.2.1"}, {"coremlc-version", "3404.23.1"}, {"coremltools-component-torch", "2.5.0"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "8.3.0"}})]
{
    func main<ios16>(tensor<fp16, [1, 1, 1, 1030]> logits) {
            tensor<int32, [1]> var_3 = const()[name = tensor<string, []>("op_3"), val = tensor<int32, [1]>([-1])];
            tensor<fp16, [1030]> flattened_cast_fp16 = reshape(shape = var_3, x = logits)[name = tensor<string, []>("flattened_cast_fp16")];
            tensor<int32, [1]> token_logits_begin_0 = const()[name = tensor<string, []>("token_logits_begin_0"), val = tensor<int32, [1]>([0])];
            tensor<int32, [1]> token_logits_end_0 = const()[name = tensor<string, []>("token_logits_end_0"), val = tensor<int32, [1]>([1025])];
            tensor<bool, [1]> token_logits_end_mask_0 = const()[name = tensor<string, []>("token_logits_end_mask_0"), val = tensor<bool, [1]>([false])];
            tensor<fp16, [1025]> token_logits_cast_fp16 = slice_by_index(begin = token_logits_begin_0, end = token_logits_end_0, end_mask = token_logits_end_mask_0, x = flattened_cast_fp16)[name = tensor<string, []>("token_logits_cast_fp16")];
            tensor<int32, [1]> duration_logits_begin_0 = const()[name = tensor<string, []>("duration_logits_begin_0"), val = tensor<int32, [1]>([1025])];
            tensor<int32, [1]> duration_logits_end_0 = const()[name = tensor<string, []>("duration_logits_end_0"), val = tensor<int32, [1]>([1])];
            tensor<bool, [1]> duration_logits_end_mask_0 = const()[name = tensor<string, []>("duration_logits_end_mask_0"), val = tensor<bool, [1]>([true])];
            tensor<fp16, [5]> duration_logits_cast_fp16 = slice_by_index(begin = duration_logits_begin_0, end = duration_logits_end_0, end_mask = duration_logits_end_mask_0, x = flattened_cast_fp16)[name = tensor<string, []>("duration_logits_cast_fp16")];
            tensor<int32, []> var_17_axis_0 = const()[name = tensor<string, []>("op_17_axis_0"), val = tensor<int32, []>(0)];
            tensor<bool, []> var_17_keep_dims_0 = const()[name = tensor<string, []>("op_17_keep_dims_0"), val = tensor<bool, []>(true)];
            tensor<int32, [1]> var_17 = reduce_argmax(axis = var_17_axis_0, keep_dims = var_17_keep_dims_0, x = token_logits_cast_fp16)[name = tensor<string, []>("op_17_cast_fp16")];
            tensor<int32, [1]> reduce_max_0_axes_0 = const()[name = tensor<string, []>("reduce_max_0_axes_0"), val = tensor<int32, [1]>([0])];
            tensor<bool, []> reduce_max_0_keep_dims_0 = const()[name = tensor<string, []>("reduce_max_0_keep_dims_0"), val = tensor<bool, []>(true)];
            tensor<fp16, [1]> reduce_max_0 = reduce_max(axes = reduce_max_0_axes_0, keep_dims = reduce_max_0_keep_dims_0, x = token_logits_cast_fp16)[name = tensor<string, []>("reduce_max_0_cast_fp16")];
            tensor<int32, []> var_24_axis_0 = const()[name = tensor<string, []>("op_24_axis_0"), val = tensor<int32, []>(0)];
            tensor<bool, []> var_24_keep_dims_0 = const()[name = tensor<string, []>("op_24_keep_dims_0"), val = tensor<bool, []>(true)];
            tensor<int32, [1]> var_24 = reduce_argmax(axis = var_24_axis_0, keep_dims = var_24_keep_dims_0, x = duration_logits_cast_fp16)[name = tensor<string, []>("op_24_cast_fp16")];
        } -> (var_17, reduce_max_0, var_24);
}