venier's picture
Upload Parakeet-TDT v2 CoreML models
6d8c66e verified
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, [2, 1, 640]> c_in, tensor<fp32, [2, 1, 640]> h_in, tensor<int32, [1]> target_length, tensor<int32, [1, 1]> targets) {
tensor<fp32, [1025, 640]> module_prediction_embed_weight = const()[name = tensor<string, []>("module_prediction_embed_weight"), val = tensor<fp32, [1025, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(64)))];
tensor<int32, []> y_batch_dims_0 = const()[name = tensor<string, []>("y_batch_dims_0"), val = tensor<int32, []>(0)];
tensor<bool, []> y_validate_indices_0 = const()[name = tensor<string, []>("y_validate_indices_0"), val = tensor<bool, []>(false)];
tensor<int32, []> greater_equal_0_y_0 = const()[name = tensor<string, []>("greater_equal_0_y_0"), val = tensor<int32, []>(0)];
tensor<bool, [1, 1]> greater_equal_0 = greater_equal(x = targets, y = greater_equal_0_y_0)[name = tensor<string, []>("greater_equal_0")];
tensor<int32, []> slice_by_index_0 = const()[name = tensor<string, []>("slice_by_index_0"), val = tensor<int32, []>(1025)];
tensor<int32, [1, 1]> add_2 = add(x = targets, y = slice_by_index_0)[name = tensor<string, []>("add_2")];
tensor<int32, [1, 1]> select_0 = select(a = targets, b = add_2, cond = greater_equal_0)[name = tensor<string, []>("select_0")];
tensor<int32, []> y_axis_1 = const()[name = tensor<string, []>("y_axis_1"), val = tensor<int32, []>(0)];
tensor<fp32, [1, 1, 640]> y = gather(axis = y_axis_1, batch_dims = y_batch_dims_0, indices = select_0, validate_indices = y_validate_indices_0, x = module_prediction_embed_weight)[name = tensor<string, []>("y")];
tensor<int32, [3]> input_3_perm_0 = const()[name = tensor<string, []>("input_3_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<fp32, [1, 1, 640]> split_0_0, tensor<fp32, [1, 1, 640]> split_0_1 = split(axis = split_0_axis_0, num_splits = split_0_num_splits_0, x = h_in)[name = tensor<string, []>("split_0")];
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<fp32, [1, 1, 640]> split_1_0, tensor<fp32, [1, 1, 640]> split_1_1 = split(axis = split_1_axis_0, num_splits = split_1_num_splits_0, x = c_in)[name = tensor<string, []>("split_1")];
tensor<fp32, [2560]> concat_0 = const()[name = tensor<string, []>("concat_0"), val = tensor<fp32, [2560]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(2624128)))];
tensor<fp32, [2560, 640]> concat_1 = const()[name = tensor<string, []>("concat_1"), val = tensor<fp32, [2560, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(2634432)))];
tensor<fp32, [2560, 640]> concat_2 = const()[name = tensor<string, []>("concat_2"), val = tensor<fp32, [2560, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(9188096)))];
tensor<int32, [1]> input_lstm_layer_0_lstm_h0_squeeze_axes_0 = const()[name = tensor<string, []>("input_lstm_layer_0_lstm_h0_squeeze_axes_0"), val = tensor<int32, [1]>([0])];
tensor<fp32, [1, 640]> input_lstm_layer_0_lstm_h0_squeeze = squeeze(axes = input_lstm_layer_0_lstm_h0_squeeze_axes_0, x = split_0_0)[name = tensor<string, []>("input_lstm_layer_0_lstm_h0_squeeze")];
tensor<int32, [1]> input_lstm_layer_0_lstm_c0_squeeze_axes_0 = const()[name = tensor<string, []>("input_lstm_layer_0_lstm_c0_squeeze_axes_0"), val = tensor<int32, [1]>([0])];
tensor<fp32, [1, 640]> input_lstm_layer_0_lstm_c0_squeeze = squeeze(axes = input_lstm_layer_0_lstm_c0_squeeze_axes_0, x = split_1_0)[name = tensor<string, []>("input_lstm_layer_0_lstm_c0_squeeze")];
tensor<string, []> input_lstm_layer_0_direction_0 = const()[name = tensor<string, []>("input_lstm_layer_0_direction_0"), val = tensor<string, []>("forward")];
tensor<bool, []> input_lstm_layer_0_output_sequence_0 = const()[name = tensor<string, []>("input_lstm_layer_0_output_sequence_0"), val = tensor<bool, []>(true)];
tensor<string, []> input_lstm_layer_0_recurrent_activation_0 = const()[name = tensor<string, []>("input_lstm_layer_0_recurrent_activation_0"), val = tensor<string, []>("sigmoid")];
tensor<string, []> input_lstm_layer_0_cell_activation_0 = const()[name = tensor<string, []>("input_lstm_layer_0_cell_activation_0"), val = tensor<string, []>("tanh")];
tensor<string, []> input_lstm_layer_0_activation_0 = const()[name = tensor<string, []>("input_lstm_layer_0_activation_0"), val = tensor<string, []>("tanh")];
tensor<fp32, [1, 1, 640]> input_3 = transpose(perm = input_3_perm_0, x = y)[name = tensor<string, []>("transpose_2")];
tensor<fp32, [1, 1, 640]> input_lstm_layer_0_0, tensor<fp32, [1, 640]> input_lstm_layer_0_1, tensor<fp32, [1, 640]> input_lstm_layer_0_2 = lstm(activation = input_lstm_layer_0_activation_0, bias = concat_0, cell_activation = input_lstm_layer_0_cell_activation_0, direction = input_lstm_layer_0_direction_0, initial_c = input_lstm_layer_0_lstm_c0_squeeze, initial_h = input_lstm_layer_0_lstm_h0_squeeze, output_sequence = input_lstm_layer_0_output_sequence_0, recurrent_activation = input_lstm_layer_0_recurrent_activation_0, weight_hh = concat_2, weight_ih = concat_1, x = input_3)[name = tensor<string, []>("input_lstm_layer_0")];
tensor<fp32, [2560]> concat_3 = const()[name = tensor<string, []>("concat_3"), val = tensor<fp32, [2560]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(15741760)))];
tensor<fp32, [2560, 640]> concat_4 = const()[name = tensor<string, []>("concat_4"), val = tensor<fp32, [2560, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(15752064)))];
tensor<fp32, [2560, 640]> concat_5 = const()[name = tensor<string, []>("concat_5"), val = tensor<fp32, [2560, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(22305728)))];
tensor<int32, [1]> input_lstm_h0_squeeze_axes_0 = const()[name = tensor<string, []>("input_lstm_h0_squeeze_axes_0"), val = tensor<int32, [1]>([0])];
tensor<fp32, [1, 640]> input_lstm_h0_squeeze = squeeze(axes = input_lstm_h0_squeeze_axes_0, x = split_0_1)[name = tensor<string, []>("input_lstm_h0_squeeze")];
tensor<int32, [1]> input_lstm_c0_squeeze_axes_0 = const()[name = tensor<string, []>("input_lstm_c0_squeeze_axes_0"), val = tensor<int32, [1]>([0])];
tensor<fp32, [1, 640]> input_lstm_c0_squeeze = squeeze(axes = input_lstm_c0_squeeze_axes_0, x = split_1_1)[name = tensor<string, []>("input_lstm_c0_squeeze")];
tensor<string, []> input_direction_0 = const()[name = tensor<string, []>("input_direction_0"), val = tensor<string, []>("forward")];
tensor<bool, []> input_output_sequence_0 = const()[name = tensor<string, []>("input_output_sequence_0"), val = tensor<bool, []>(true)];
tensor<string, []> input_recurrent_activation_0 = const()[name = tensor<string, []>("input_recurrent_activation_0"), val = tensor<string, []>("sigmoid")];
tensor<string, []> input_cell_activation_0 = const()[name = tensor<string, []>("input_cell_activation_0"), val = tensor<string, []>("tanh")];
tensor<string, []> input_activation_0 = const()[name = tensor<string, []>("input_activation_0"), val = tensor<string, []>("tanh")];
tensor<fp32, [1, 1, 640]> input_0, tensor<fp32, [1, 640]> input_1, tensor<fp32, [1, 640]> input_2 = lstm(activation = input_activation_0, bias = concat_3, cell_activation = input_cell_activation_0, direction = input_direction_0, initial_c = input_lstm_c0_squeeze, initial_h = input_lstm_h0_squeeze, output_sequence = input_output_sequence_0, recurrent_activation = input_recurrent_activation_0, weight_hh = concat_5, weight_ih = concat_4, x = input_lstm_layer_0_0)[name = tensor<string, []>("input")];
tensor<int32, []> obj_3_axis_0 = const()[name = tensor<string, []>("obj_3_axis_0"), val = tensor<int32, []>(0)];
tensor<fp32, [2, 1, 640]> h_out = stack(axis = obj_3_axis_0, values = (input_lstm_layer_0_1, input_1))[name = tensor<string, []>("obj_3")];
tensor<int32, []> obj_axis_0 = const()[name = tensor<string, []>("obj_axis_0"), val = tensor<int32, []>(0)];
tensor<fp32, [2, 1, 640]> c_out = stack(axis = obj_axis_0, values = (input_lstm_layer_0_2, input_2))[name = tensor<string, []>("obj")];
tensor<int32, [3]> transpose_0_perm_0 = const()[name = tensor<string, []>("transpose_0_perm_0"), val = tensor<int32, [3]>([1, 2, 0])];
tensor<fp32, [1, 640, 1]> decoder = transpose(perm = transpose_0_perm_0, x = input_0)[name = tensor<string, []>("transpose_1")];
tensor<int32, [1]> target_length_tmp = identity(x = target_length)[name = tensor<string, []>("target_length_tmp")];
} -> (decoder, h_out, c_out);
}