program(1.3)
[buildInfo = dict<string, string>({{"coremlc-component-MIL", "3510.2.1"}, {"coremlc-version", "3500.32.1"}, {"coremltools-component-torch", "2.11.0"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "9.0"}})]
{
    func main<ios18>(tensor<fp16, [1, 2048, 1, 1]> lm_hidden, tensor<fp16, [1, 2048, 1, 1]> residual_hidden) {
            tensor<int32, [1]> var_8_axes_0 = const()[name = string("op_8_axes_0"), val = tensor<int32, [1]>([-1])];
            tensor<fp16, [1, 2048, 1]> var_8_cast_fp16 = squeeze(axes = var_8_axes_0, x = lm_hidden)[name = string("op_8_cast_fp16")];
            tensor<int32, [1]> input_1_axes_0 = const()[name = string("input_1_axes_0"), val = tensor<int32, [1]>([-1])];
            tensor<fp16, [1, 2048]> input_1_cast_fp16 = squeeze(axes = input_1_axes_0, x = var_8_cast_fp16)[name = string("input_1_cast_fp16")];
            tensor<int32, [1]> var_12_axes_0 = const()[name = string("op_12_axes_0"), val = tensor<int32, [1]>([-1])];
            tensor<fp16, [1, 2048, 1]> var_12_cast_fp16 = squeeze(axes = var_12_axes_0, x = residual_hidden)[name = string("op_12_cast_fp16")];
            tensor<int32, [1]> input_3_axes_0 = const()[name = string("input_3_axes_0"), val = tensor<int32, [1]>([-1])];
            tensor<fp16, [1, 2048]> input_3_cast_fp16 = squeeze(axes = input_3_axes_0, x = var_12_cast_fp16)[name = string("input_3_cast_fp16")];
            tensor<fp16, [1024, 2048]> lm_to_dit_proj_weight_to_fp16 = const()[name = string("lm_to_dit_proj_weight_to_fp16"), val = tensor<fp16, [1024, 2048]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(64)))];
            tensor<fp16, [1024]> lm_to_dit_proj_bias_to_fp16 = const()[name = string("lm_to_dit_proj_bias_to_fp16"), val = tensor<fp16, [1024]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(4194432)))];
            tensor<fp16, [1, 1024]> linear_0_cast_fp16 = linear(bias = lm_to_dit_proj_bias_to_fp16, weight = lm_to_dit_proj_weight_to_fp16, x = input_1_cast_fp16)[name = string("linear_0_cast_fp16")];
            tensor<fp16, [1024, 2048]> res_to_dit_proj_weight_to_fp16 = const()[name = string("res_to_dit_proj_weight_to_fp16"), val = tensor<fp16, [1024, 2048]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(4196544)))];
            tensor<fp16, [1024]> res_to_dit_proj_bias_to_fp16 = const()[name = string("res_to_dit_proj_bias_to_fp16"), val = tensor<fp16, [1024]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(8390912)))];
            tensor<fp16, [1, 1024]> linear_1_cast_fp16 = linear(bias = res_to_dit_proj_bias_to_fp16, weight = res_to_dit_proj_weight_to_fp16, x = input_3_cast_fp16)[name = string("linear_1_cast_fp16")];
            int32 var_22 = const()[name = string("op_22"), val = int32(-1)];
            bool dit_hidden_interleave_0 = const()[name = string("dit_hidden_interleave_0"), val = bool(false)];
            tensor<fp16, [1, 2048]> dit_hidden_cast_fp16 = concat(axis = var_22, interleave = dit_hidden_interleave_0, values = (linear_0_cast_fp16, linear_1_cast_fp16))[name = string("dit_hidden_cast_fp16")];
            tensor<fp16, [2048, 2048]> stop_proj_weight_to_fp16 = const()[name = string("stop_proj_weight_to_fp16"), val = tensor<fp16, [2048, 2048]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(8393024)))];
            tensor<fp16, [2048]> stop_proj_bias_to_fp16 = const()[name = string("stop_proj_bias_to_fp16"), val = tensor<fp16, [2048]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(16781696)))];
            tensor<fp16, [1, 2048]> linear_2_cast_fp16 = linear(bias = stop_proj_bias_to_fp16, weight = stop_proj_weight_to_fp16, x = input_1_cast_fp16)[name = string("linear_2_cast_fp16")];
            tensor<fp16, [1, 2048]> input_cast_fp16 = silu(x = linear_2_cast_fp16)[name = string("input_cast_fp16")];
            tensor<fp16, [2, 2048]> stop_head_weight_to_fp16 = const()[name = string("stop_head_weight_to_fp16"), val = tensor<fp16, [2, 2048]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(16785856)))];
            tensor<fp16, [2]> linear_3_bias_0_to_fp16 = const()[name = string("linear_3_bias_0_to_fp16"), val = tensor<fp16, [2]>([0x0p+0, 0x0p+0])];
            tensor<fp16, [1, 2]> linear_3_cast_fp16 = linear(bias = linear_3_bias_0_to_fp16, weight = stop_head_weight_to_fp16, x = input_cast_fp16)[name = string("linear_3_cast_fp16")];
            tensor<int32, [1]> var_32_axes_0 = const()[name = string("op_32_axes_0"), val = tensor<int32, [1]>([-1])];
            tensor<fp16, [1, 2048, 1]> var_32_cast_fp16 = expand_dims(axes = var_32_axes_0, x = dit_hidden_cast_fp16)[name = string("op_32_cast_fp16")];
            tensor<int32, [1]> var_34_axes_0 = const()[name = string("op_34_axes_0"), val = tensor<int32, [1]>([-1])];
            tensor<fp16, [1, 2048, 1, 1]> dit_hidden = expand_dims(axes = var_34_axes_0, x = var_32_cast_fp16)[name = string("op_34_cast_fp16")];
            tensor<int32, [1]> var_36_axes_0 = const()[name = string("op_36_axes_0"), val = tensor<int32, [1]>([-1])];
            tensor<fp16, [1, 2, 1]> var_36_cast_fp16 = expand_dims(axes = var_36_axes_0, x = linear_3_cast_fp16)[name = string("op_36_cast_fp16")];
            tensor<int32, [1]> var_38_axes_0 = const()[name = string("op_38_axes_0"), val = tensor<int32, [1]>([-1])];
            tensor<fp16, [1, 2, 1, 1]> stop_flag = expand_dims(axes = var_38_axes_0, x = var_36_cast_fp16)[name = string("op_38_cast_fp16")];
        } -> (dit_hidden, stop_flag);
}