program(1.0)
[buildInfo = dict<tensor<string, []>, tensor<string, []>>({{"coremlc-component-MIL", "3401.3.1"}, {"coremlc-version", "3401.4.1"}, {"coremltools-component-torch", "2.6.0"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "8.2"}})]
{
    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_axis_0 = const()[name = tensor<string, []>("input_1_axis_0"), val = tensor<int32, []>(0)];
            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, [1025, 640]> prediction_embed_weight_to_fp16 = const()[name = tensor<string, []>("prediction_embed_weight_to_fp16"), val = tensor<fp16, [1025, 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<int16, [1]> decoder_input_ids_to_int16 = cast(dtype = decoder_input_ids_to_int16_dtype_0, x = decoder_input_ids)[name = tensor<string, []>("cast_6")];
            tensor<fp16, [1, 640]> input_1_cast_fp16_cast_uint16 = gather(axis = input_1_axis_0, batch_dims = input_1_batch_dims_0, indices = decoder_input_ids_to_int16, validate_indices = input_1_validate_indices_0, x = prediction_embed_weight_to_fp16)[name = tensor<string, []>("input_1_cast_fp16_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)[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, []>(1312128)))];
            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, []>(4588992)))];
            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, []>(7865856)))];
            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, []>(7871040)))];
            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, []>(11147904)))];
            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, []>(14424768)))];
            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, []>(14429952)))];
            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, []>(15249216)))];
            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);
}