| program(1.0) | |
| [buildInfo = dict<tensor<string, []>, tensor<string, []>>({{"coremlc-component-MIL", "3401.3.1"}, {"coremlc-version", "3401.4.1"}, {"coremltools-component-torch", "2.5.0"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "9.0b1"}})] | |
| { | |
| func main<ios17>(tensor<int32, [1]> decoder_input_ids, tensor<fp16, [2, 640]> state_1, tensor<fp16, [2, 640]> state_2) { | |
| tensor<int32, []> input_1_batch_dims_0 = const()[name = tensor<string, []>("input_1_batch_dims_0"), val = tensor<int32, []>(0)]; | |
| tensor<bool, []> input_1_validate_indices_0 = const()[name = tensor<string, []>("input_1_validate_indices_0"), val = tensor<bool, []>(false)]; | |
| tensor<fp16, [8193, 640]> prediction_embed_weight_to_fp16 = const()[name = tensor<string, []>("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, []> decoder_input_ids_to_int16_dtype_0 = const()[name = tensor<string, []>("decoder_input_ids_to_int16_dtype_0"), val = tensor<string, []>("int16")]; | |
| tensor<string, []> cast_6_dtype_0 = const()[name = tensor<string, []>("cast_6_dtype_0"), val = tensor<string, []>("int32")]; | |
| tensor<int32, []> greater_equal_0_y_0 = const()[name = tensor<string, []>("greater_equal_0_y_0"), val = tensor<int32, []>(0)]; | |
| tensor<int16, [1]> decoder_input_ids_to_int16 = cast(dtype = decoder_input_ids_to_int16_dtype_0, x = decoder_input_ids)[name = tensor<string, []>("cast_9")]; | |
| tensor<int32, [1]> cast_6 = cast(dtype = cast_6_dtype_0, x = decoder_input_ids_to_int16)[name = tensor<string, []>("cast_8")]; | |
| tensor<bool, [1]> greater_equal_0 = greater_equal(x = cast_6, 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, []>(8193)]; | |
| tensor<int32, [1]> add_2 = add(x = cast_6, y = slice_by_index_0)[name = tensor<string, []>("add_2")]; | |
| tensor<int32, [1]> select_0 = select(a = cast_6, b = add_2, cond = greater_equal_0)[name = tensor<string, []>("select_0")]; | |
| tensor<int32, []> input_1_cast_fp16_cast_uint16_axis_0 = const()[name = tensor<string, []>("input_1_cast_fp16_cast_uint16_axis_0"), val = tensor<int32, []>(0)]; | |
| tensor<string, []> select_0_to_int16_dtype_0 = const()[name = tensor<string, []>("select_0_to_int16_dtype_0"), val = tensor<string, []>("int16")]; | |
| tensor<int16, [1]> select_0_to_int16 = cast(dtype = select_0_to_int16_dtype_0, x = select_0)[name = tensor<string, []>("cast_7")]; | |
| tensor<fp16, [1, 640]> input_1_cast_fp16_cast_uint16_cast_uint16 = gather(axis = input_1_cast_fp16_cast_uint16_axis_0, batch_dims = input_1_batch_dims_0, indices = select_0_to_int16, validate_indices = input_1_validate_indices_0, x = prediction_embed_weight_to_fp16)[name = tensor<string, []>("input_1_cast_fp16_cast_uint16_cast_uint16")]; | |
| tensor<int32, [1]> input_3_axes_0 = const()[name = tensor<string, []>("input_3_axes_0"), val = tensor<int32, [1]>([1])]; | |
| tensor<fp16, [1, 1, 640]> input_3_cast_fp16 = expand_dims(axes = input_3_axes_0, x = input_1_cast_fp16_cast_uint16_cast_uint16)[name = tensor<string, []>("input_3_cast_fp16")]; | |
| tensor<int32, [1]> hx_1_axes_0 = const()[name = tensor<string, []>("hx_1_axes_0"), val = tensor<int32, [1]>([1])]; | |
| tensor<fp16, [2, 1, 640]> hx_1_cast_fp16 = expand_dims(axes = hx_1_axes_0, x = state_1)[name = tensor<string, []>("hx_1_cast_fp16")]; | |
| tensor<int32, [1]> hx_axes_0 = const()[name = tensor<string, []>("hx_axes_0"), val = tensor<int32, [1]>([1])]; | |
| tensor<fp16, [2, 1, 640]> hx_cast_fp16 = expand_dims(axes = hx_axes_0, x = state_2)[name = tensor<string, []>("hx_cast_fp16")]; | |
| 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 = hx_1_cast_fp16)[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 = hx_cast_fp16)[name = tensor<string, []>("split_1_cast_fp16")]; | |
| tensor<int32, [1]> output_lstm_layer_0_lstm_h0_squeeze_axes_0 = const()[name = tensor<string, []>("output_lstm_layer_0_lstm_h0_squeeze_axes_0"), val = tensor<int32, [1]>([0])]; | |
| tensor<fp16, [1, 640]> output_lstm_layer_0_lstm_h0_squeeze_cast_fp16 = squeeze(axes = output_lstm_layer_0_lstm_h0_squeeze_axes_0, x = split_0_cast_fp16_0)[name = tensor<string, []>("output_lstm_layer_0_lstm_h0_squeeze_cast_fp16")]; | |
| tensor<int32, [1]> output_lstm_layer_0_lstm_c0_squeeze_axes_0 = const()[name = tensor<string, []>("output_lstm_layer_0_lstm_c0_squeeze_axes_0"), val = tensor<int32, [1]>([0])]; | |
| tensor<fp16, [1, 640]> output_lstm_layer_0_lstm_c0_squeeze_cast_fp16 = squeeze(axes = output_lstm_layer_0_lstm_c0_squeeze_axes_0, x = split_1_cast_fp16_0)[name = tensor<string, []>("output_lstm_layer_0_lstm_c0_squeeze_cast_fp16")]; | |
| tensor<string, []> output_lstm_layer_0_direction_0 = const()[name = tensor<string, []>("output_lstm_layer_0_direction_0"), val = tensor<string, []>("forward")]; | |
| tensor<bool, []> output_lstm_layer_0_output_sequence_0 = const()[name = tensor<string, []>("output_lstm_layer_0_output_sequence_0"), val = tensor<bool, []>(true)]; | |
| tensor<string, []> output_lstm_layer_0_recurrent_activation_0 = const()[name = tensor<string, []>("output_lstm_layer_0_recurrent_activation_0"), val = tensor<string, []>("sigmoid")]; | |
| tensor<string, []> output_lstm_layer_0_cell_activation_0 = const()[name = tensor<string, []>("output_lstm_layer_0_cell_activation_0"), val = tensor<string, []>("tanh")]; | |
| tensor<string, []> output_lstm_layer_0_activation_0 = const()[name = tensor<string, []>("output_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]> output_lstm_layer_0_cast_fp16_0, tensor<fp16, [1, 640]> output_lstm_layer_0_cast_fp16_1, tensor<fp16, [1, 640]> output_lstm_layer_0_cast_fp16_2 = lstm(activation = output_lstm_layer_0_activation_0, bias = concat_0_to_fp16, cell_activation = output_lstm_layer_0_cell_activation_0, direction = output_lstm_layer_0_direction_0, initial_c = output_lstm_layer_0_lstm_c0_squeeze_cast_fp16, initial_h = output_lstm_layer_0_lstm_h0_squeeze_cast_fp16, output_sequence = output_lstm_layer_0_output_sequence_0, recurrent_activation = output_lstm_layer_0_recurrent_activation_0, weight_hh = concat_2_to_fp16, weight_ih = concat_1_to_fp16, x = input_3_cast_fp16)[name = tensor<string, []>("output_lstm_layer_0_cast_fp16")]; | |
| tensor<int32, [1]> output_lstm_h0_squeeze_axes_0 = const()[name = tensor<string, []>("output_lstm_h0_squeeze_axes_0"), val = tensor<int32, [1]>([0])]; | |
| tensor<fp16, [1, 640]> output_lstm_h0_squeeze_cast_fp16 = squeeze(axes = output_lstm_h0_squeeze_axes_0, x = split_0_cast_fp16_1)[name = tensor<string, []>("output_lstm_h0_squeeze_cast_fp16")]; | |
| tensor<int32, [1]> output_lstm_c0_squeeze_axes_0 = const()[name = tensor<string, []>("output_lstm_c0_squeeze_axes_0"), val = tensor<int32, [1]>([0])]; | |
| tensor<fp16, [1, 640]> output_lstm_c0_squeeze_cast_fp16 = squeeze(axes = output_lstm_c0_squeeze_axes_0, x = split_1_cast_fp16_1)[name = tensor<string, []>("output_lstm_c0_squeeze_cast_fp16")]; | |
| tensor<string, []> output_direction_0 = const()[name = tensor<string, []>("output_direction_0"), val = tensor<string, []>("forward")]; | |
| tensor<bool, []> output_output_sequence_0 = const()[name = tensor<string, []>("output_output_sequence_0"), val = tensor<bool, []>(true)]; | |
| tensor<string, []> output_recurrent_activation_0 = const()[name = tensor<string, []>("output_recurrent_activation_0"), val = tensor<string, []>("sigmoid")]; | |
| tensor<string, []> output_cell_activation_0 = const()[name = tensor<string, []>("output_cell_activation_0"), val = tensor<string, []>("tanh")]; | |
| tensor<string, []> output_activation_0 = const()[name = tensor<string, []>("output_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]> output_cast_fp16_0, tensor<fp16, [1, 640]> output_cast_fp16_1, tensor<fp16, [1, 640]> output_cast_fp16_2 = lstm(activation = output_activation_0, bias = concat_3_to_fp16, cell_activation = output_cell_activation_0, direction = output_direction_0, initial_c = output_lstm_c0_squeeze_cast_fp16, initial_h = output_lstm_h0_squeeze_cast_fp16, output_sequence = output_output_sequence_0, recurrent_activation = output_recurrent_activation_0, weight_hh = concat_5_to_fp16, weight_ih = concat_4_to_fp16, x = output_lstm_layer_0_cast_fp16_0)[name = tensor<string, []>("output_cast_fp16")]; | |
| tensor<int32, []> var_32_axis_0 = const()[name = tensor<string, []>("op_32_axis_0"), val = tensor<int32, []>(0)]; | |
| tensor<fp16, [2, 1, 640]> var_32_cast_fp16 = stack(axis = var_32_axis_0, values = (output_lstm_layer_0_cast_fp16_1, output_cast_fp16_1))[name = tensor<string, []>("op_32_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]> var_33_cast_fp16 = stack(axis = var_33_axis_0, values = (output_lstm_layer_0_cast_fp16_2, output_cast_fp16_2))[name = tensor<string, []>("op_33_cast_fp16")]; | |
| tensor<int32, [1]> input_axes_0 = const()[name = tensor<string, []>("input_axes_0"), val = tensor<int32, [1]>([1])]; | |
| tensor<fp16, [1, 640]> input_cast_fp16 = squeeze(axes = input_axes_0, x = output_cast_fp16_0)[name = tensor<string, []>("input_cast_fp16")]; | |
| tensor<int32, [1]> var_35_axes_0 = const()[name = tensor<string, []>("op_35_axes_0"), val = tensor<int32, [1]>([1])]; | |
| tensor<fp16, [2, 640]> new_state_1 = squeeze(axes = var_35_axes_0, x = var_32_cast_fp16)[name = tensor<string, []>("op_35_cast_fp16")]; | |
| tensor<int32, [1]> var_36_axes_0 = const()[name = tensor<string, []>("op_36_axes_0"), val = tensor<int32, [1]>([1])]; | |
| tensor<fp16, [2, 640]> new_state_2 = squeeze(axes = var_36_axes_0, x = var_33_cast_fp16)[name = tensor<string, []>("op_36_cast_fp16")]; | |
| tensor<fp16, [640, 640]> joint_projection_weight_to_fp16 = const()[name = tensor<string, []>("joint_projection_weight_to_fp16"), val = tensor<fp16, [640, 640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(23604992)))]; | |
| tensor<fp16, [640]> joint_projection_bias_to_fp16 = const()[name = tensor<string, []>("joint_projection_bias_to_fp16"), val = tensor<fp16, [640]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(24424256)))]; | |
| tensor<fp16, [1, 640]> decoder_output_projected = linear(bias = joint_projection_bias_to_fp16, weight = joint_projection_weight_to_fp16, x = input_cast_fp16)[name = tensor<string, []>("linear_0_cast_fp16")]; | |
| } -> (decoder_output_projected, new_state_1, new_state_2); | |
| } |