| program(1.3) |
| [buildInfo = dict<string, string>({{"coremlc-component-MIL", "3600.16.1"}, {"coremlc-version", "3600.25.2"}, {"coremltools-component-torch", "2.11.0"}, {"coremltools-source-dialect", "TorchExport::ATEN"}, {"coremltools-version", "9.0"}})] |
| { |
| func main<ios18>(tensor<fp32, [2, 1, 640]> c_in, tensor<fp32, [2, 1, 640]> h_in, tensor<int32, [1, 1]> token) { |
| int32 embedding_batch_dims_0 = const()[name = string("embedding_batch_dims_0"), val = int32(0)]; |
| bool embedding_validate_indices_0 = const()[name = string("embedding_validate_indices_0"), val = bool(false)]; |
| tensor<fp16, [8193, 640]> p_embedding_weight_to_fp16 = const()[name = string("p_embedding_weight_to_fp16"), val = tensor<fp16, [8193, 640]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(64)))]; |
| string token_to_int16_dtype_0 = const()[name = string("token_to_int16_dtype_0"), val = string("int16")]; |
| string cast_4_dtype_0 = const()[name = string("cast_4_dtype_0"), val = string("int32")]; |
| int32 greater_equal_0_y_0 = const()[name = string("greater_equal_0_y_0"), val = int32(0)]; |
| tensor<int16, [1, 1]> token_to_int16 = cast(dtype = token_to_int16_dtype_0, x = token)[name = string("cast_9")]; |
| tensor<int32, [1, 1]> cast_4 = cast(dtype = cast_4_dtype_0, x = token_to_int16)[name = string("cast_8")]; |
| tensor<bool, [1, 1]> greater_equal_0 = greater_equal(x = cast_4, y = greater_equal_0_y_0)[name = string("greater_equal_0")]; |
| int32 slice_by_index_0 = const()[name = string("slice_by_index_0"), val = int32(8193)]; |
| tensor<int32, [1, 1]> add_0 = add(x = cast_4, y = slice_by_index_0)[name = string("add_0")]; |
| tensor<int32, [1, 1]> select_0 = select(a = cast_4, b = add_0, cond = greater_equal_0)[name = string("select_0")]; |
| int32 embedding_cast_fp16_cast_uint16_axis_0 = const()[name = string("embedding_cast_fp16_cast_uint16_axis_0"), val = int32(0)]; |
| string select_0_to_int16_dtype_0 = const()[name = string("select_0_to_int16_dtype_0"), val = string("int16")]; |
| tensor<int16, [1, 1]> select_0_to_int16 = cast(dtype = select_0_to_int16_dtype_0, x = select_0)[name = string("cast_7")]; |
| tensor<fp16, [1, 1, 640]> embedding_cast_fp16_cast_uint16_cast_uint16 = gather(axis = embedding_cast_fp16_cast_uint16_axis_0, batch_dims = embedding_batch_dims_0, indices = select_0_to_int16, validate_indices = embedding_validate_indices_0, x = p_embedding_weight_to_fp16)[name = string("embedding_cast_fp16_cast_uint16_cast_uint16")]; |
| tensor<int32, [3]> select_begin_0 = const()[name = string("select_begin_0"), val = tensor<int32, [3]>([0, 0, 0])]; |
| tensor<int32, [3]> select_end_0 = const()[name = string("select_end_0"), val = tensor<int32, [3]>([1, 1, 640])]; |
| tensor<bool, [3]> select_end_mask_0 = const()[name = string("select_end_mask_0"), val = tensor<bool, [3]>([true, false, true])]; |
| tensor<bool, [3]> select_squeeze_mask_0 = const()[name = string("select_squeeze_mask_0"), val = tensor<bool, [3]>([false, true, false])]; |
| tensor<fp16, [1, 640]> select_cast_fp16 = slice_by_index(begin = select_begin_0, end = select_end_0, end_mask = select_end_mask_0, squeeze_mask = select_squeeze_mask_0, x = embedding_cast_fp16_cast_uint16_cast_uint16)[name = string("select_cast_fp16")]; |
| tensor<int32, [3]> select_1_begin_0 = const()[name = string("select_1_begin_0"), val = tensor<int32, [3]>([0, 0, 0])]; |
| tensor<int32, [3]> select_1_end_0 = const()[name = string("select_1_end_0"), val = tensor<int32, [3]>([1, 1, 640])]; |
| tensor<bool, [3]> select_1_end_mask_0 = const()[name = string("select_1_end_mask_0"), val = tensor<bool, [3]>([false, true, true])]; |
| tensor<bool, [3]> select_1_squeeze_mask_0 = const()[name = string("select_1_squeeze_mask_0"), val = tensor<bool, [3]>([true, false, false])]; |
| string h_in_to_fp16_dtype_0 = const()[name = string("h_in_to_fp16_dtype_0"), val = string("fp16")]; |
| tensor<fp16, [2, 1, 640]> h_in_to_fp16 = cast(dtype = h_in_to_fp16_dtype_0, x = h_in)[name = string("cast_6")]; |
| tensor<fp16, [1, 640]> select_1_cast_fp16 = slice_by_index(begin = select_1_begin_0, end = select_1_end_0, end_mask = select_1_end_mask_0, squeeze_mask = select_1_squeeze_mask_0, x = h_in_to_fp16)[name = string("select_1_cast_fp16")]; |
| tensor<int32, [3]> select_2_begin_0 = const()[name = string("select_2_begin_0"), val = tensor<int32, [3]>([0, 0, 0])]; |
| tensor<int32, [3]> select_2_end_0 = const()[name = string("select_2_end_0"), val = tensor<int32, [3]>([1, 1, 640])]; |
| tensor<bool, [3]> select_2_end_mask_0 = const()[name = string("select_2_end_mask_0"), val = tensor<bool, [3]>([false, true, true])]; |
| tensor<bool, [3]> select_2_squeeze_mask_0 = const()[name = string("select_2_squeeze_mask_0"), val = tensor<bool, [3]>([true, false, false])]; |
| string c_in_to_fp16_dtype_0 = const()[name = string("c_in_to_fp16_dtype_0"), val = string("fp16")]; |
| tensor<fp16, [2, 1, 640]> c_in_to_fp16 = cast(dtype = c_in_to_fp16_dtype_0, x = c_in)[name = string("cast_5")]; |
| tensor<fp16, [1, 640]> select_2_cast_fp16 = slice_by_index(begin = select_2_begin_0, end = select_2_end_0, end_mask = select_2_end_mask_0, squeeze_mask = select_2_squeeze_mask_0, x = c_in_to_fp16)[name = string("select_2_cast_fp16")]; |
| tensor<fp16, [2560, 640]> add_weight_0_to_fp16 = const()[name = string("add_weight_0_to_fp16"), val = tensor<fp16, [2560, 640]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(10487168)))]; |
| tensor<fp16, [2560]> add_bias_0_to_fp16 = const()[name = string("add_bias_0_to_fp16"), val = tensor<fp16, [2560]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(13764032)))]; |
| tensor<fp16, [1, 2560]> add_cast_fp16 = linear(bias = add_bias_0_to_fp16, weight = add_weight_0_to_fp16, x = select_cast_fp16)[name = string("add_cast_fp16")]; |
| tensor<fp16, [2560, 640]> b_weight_hh_l0_to_fp16 = const()[name = string("b_weight_hh_l0_to_fp16"), val = tensor<fp16, [2560, 640]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(13769216)))]; |
| tensor<fp16, [2560]> matmul_bias_0_to_fp16 = const()[name = string("matmul_bias_0_to_fp16"), val = tensor<fp16, [2560]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(17046080)))]; |
| tensor<fp16, [1, 2560]> matmul_1_cast_fp16 = linear(bias = matmul_bias_0_to_fp16, weight = b_weight_hh_l0_to_fp16, x = select_1_cast_fp16)[name = string("matmul_1_cast_fp16")]; |
| tensor<fp16, [1, 2560]> add_1_cast_fp16 = add(x = add_cast_fp16, y = matmul_1_cast_fp16)[name = string("add_1_cast_fp16")]; |
| tensor<fp16, [2560]> b_bias_hh_l0_to_fp16 = const()[name = string("b_bias_hh_l0_to_fp16"), val = tensor<fp16, [2560]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(17051264)))]; |
| tensor<fp16, [1, 2560]> add_2_cast_fp16 = add(x = add_1_cast_fp16, y = b_bias_hh_l0_to_fp16)[name = string("add_2_cast_fp16")]; |
| tensor<int32, [4]> tile_0 = const()[name = string("tile_0"), val = tensor<int32, [4]>([640, 640, 640, 640])]; |
| int32 split_axis_0 = const()[name = string("split_axis_0"), val = int32(-1)]; |
| tensor<fp16, [1, 640]> split_cast_fp16_0, tensor<fp16, [1, 640]> split_cast_fp16_1, tensor<fp16, [1, 640]> split_cast_fp16_2, tensor<fp16, [1, 640]> split_cast_fp16_3 = split(axis = split_axis_0, split_sizes = tile_0, x = add_2_cast_fp16)[name = string("split_cast_fp16")]; |
| tensor<fp16, [1, 640]> sigmoid_cast_fp16 = sigmoid(x = split_cast_fp16_1)[name = string("sigmoid_cast_fp16")]; |
| tensor<fp16, [1, 640]> mul_cast_fp16 = mul(x = sigmoid_cast_fp16, y = select_2_cast_fp16)[name = string("mul_cast_fp16")]; |
| tensor<fp16, [1, 640]> sigmoid_1_cast_fp16 = sigmoid(x = split_cast_fp16_0)[name = string("sigmoid_1_cast_fp16")]; |
| tensor<fp16, [1, 640]> tanh_cast_fp16 = tanh(x = split_cast_fp16_2)[name = string("tanh_cast_fp16")]; |
| tensor<fp16, [1, 640]> mul_1_1_cast_fp16 = mul(x = sigmoid_1_cast_fp16, y = tanh_cast_fp16)[name = string("mul_1_1_cast_fp16")]; |
| tensor<fp16, [1, 640]> add_3_cast_fp16 = add(x = mul_cast_fp16, y = mul_1_1_cast_fp16)[name = string("add_3_cast_fp16")]; |
| tensor<fp16, [1, 640]> sigmoid_2_cast_fp16 = sigmoid(x = split_cast_fp16_3)[name = string("sigmoid_2_cast_fp16")]; |
| tensor<fp16, [1, 640]> tanh_1_cast_fp16 = tanh(x = add_3_cast_fp16)[name = string("tanh_1_cast_fp16")]; |
| tensor<fp16, [1, 640]> mul_2_cast_fp16 = mul(x = sigmoid_2_cast_fp16, y = tanh_1_cast_fp16)[name = string("mul_2_cast_fp16")]; |
| tensor<int32, [3]> select_3_begin_0 = const()[name = string("select_3_begin_0"), val = tensor<int32, [3]>([1, 0, 0])]; |
| tensor<int32, [3]> select_3_end_0 = const()[name = string("select_3_end_0"), val = tensor<int32, [3]>([2, 1, 640])]; |
| tensor<bool, [3]> select_3_end_mask_0 = const()[name = string("select_3_end_mask_0"), val = tensor<bool, [3]>([false, true, true])]; |
| tensor<bool, [3]> select_3_squeeze_mask_0 = const()[name = string("select_3_squeeze_mask_0"), val = tensor<bool, [3]>([true, false, false])]; |
| tensor<fp16, [1, 640]> select_3_cast_fp16 = slice_by_index(begin = select_3_begin_0, end = select_3_end_0, end_mask = select_3_end_mask_0, squeeze_mask = select_3_squeeze_mask_0, x = h_in_to_fp16)[name = string("select_3_cast_fp16")]; |
| tensor<int32, [3]> select_4_begin_0 = const()[name = string("select_4_begin_0"), val = tensor<int32, [3]>([1, 0, 0])]; |
| tensor<int32, [3]> select_4_end_0 = const()[name = string("select_4_end_0"), val = tensor<int32, [3]>([2, 1, 640])]; |
| tensor<bool, [3]> select_4_end_mask_0 = const()[name = string("select_4_end_mask_0"), val = tensor<bool, [3]>([false, true, true])]; |
| tensor<bool, [3]> select_4_squeeze_mask_0 = const()[name = string("select_4_squeeze_mask_0"), val = tensor<bool, [3]>([true, false, false])]; |
| tensor<fp16, [1, 640]> select_4_cast_fp16 = slice_by_index(begin = select_4_begin_0, end = select_4_end_0, end_mask = select_4_end_mask_0, squeeze_mask = select_4_squeeze_mask_0, x = c_in_to_fp16)[name = string("select_4_cast_fp16")]; |
| tensor<fp16, [2560, 640]> add_4_weight_0_to_fp16 = const()[name = string("add_4_weight_0_to_fp16"), val = tensor<fp16, [2560, 640]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(17056448)))]; |
| tensor<fp16, [2560]> add_4_bias_0_to_fp16 = const()[name = string("add_4_bias_0_to_fp16"), val = tensor<fp16, [2560]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(20333312)))]; |
| tensor<fp16, [1, 2560]> add_4_cast_fp16 = linear(bias = add_4_bias_0_to_fp16, weight = add_4_weight_0_to_fp16, x = mul_2_cast_fp16)[name = string("add_4_cast_fp16")]; |
| tensor<fp16, [2560, 640]> b_weight_hh_l1_to_fp16 = const()[name = string("b_weight_hh_l1_to_fp16"), val = tensor<fp16, [2560, 640]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(20338496)))]; |
| tensor<fp16, [1, 2560]> matmul_3_cast_fp16 = linear(bias = matmul_bias_0_to_fp16, weight = b_weight_hh_l1_to_fp16, x = select_3_cast_fp16)[name = string("matmul_3_cast_fp16")]; |
| tensor<fp16, [1, 2560]> add_5_cast_fp16 = add(x = add_4_cast_fp16, y = matmul_3_cast_fp16)[name = string("add_5_cast_fp16")]; |
| tensor<fp16, [2560]> b_bias_hh_l1_to_fp16 = const()[name = string("b_bias_hh_l1_to_fp16"), val = tensor<fp16, [2560]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(23615360)))]; |
| tensor<fp16, [1, 2560]> add_6_cast_fp16 = add(x = add_5_cast_fp16, y = b_bias_hh_l1_to_fp16)[name = string("add_6_cast_fp16")]; |
| tensor<int32, [4]> tile_1 = const()[name = string("tile_1"), val = tensor<int32, [4]>([640, 640, 640, 640])]; |
| int32 split_1_axis_0 = const()[name = string("split_1_axis_0"), val = int32(-1)]; |
| tensor<fp16, [1, 640]> split_1_cast_fp16_0, tensor<fp16, [1, 640]> split_1_cast_fp16_1, tensor<fp16, [1, 640]> split_1_cast_fp16_2, tensor<fp16, [1, 640]> split_1_cast_fp16_3 = split(axis = split_1_axis_0, split_sizes = tile_1, x = add_6_cast_fp16)[name = string("split_1_cast_fp16")]; |
| tensor<fp16, [1, 640]> sigmoid_3_cast_fp16 = sigmoid(x = split_1_cast_fp16_1)[name = string("sigmoid_3_cast_fp16")]; |
| tensor<fp16, [1, 640]> mul_3_1_cast_fp16 = mul(x = sigmoid_3_cast_fp16, y = select_4_cast_fp16)[name = string("mul_3_1_cast_fp16")]; |
| tensor<fp16, [1, 640]> sigmoid_4_cast_fp16 = sigmoid(x = split_1_cast_fp16_0)[name = string("sigmoid_4_cast_fp16")]; |
| tensor<fp16, [1, 640]> tanh_2_cast_fp16 = tanh(x = split_1_cast_fp16_2)[name = string("tanh_2_cast_fp16")]; |
| tensor<fp16, [1, 640]> mul_4_cast_fp16 = mul(x = sigmoid_4_cast_fp16, y = tanh_2_cast_fp16)[name = string("mul_4_cast_fp16")]; |
| tensor<fp16, [1, 640]> add_7_cast_fp16 = add(x = mul_3_1_cast_fp16, y = mul_4_cast_fp16)[name = string("add_7_cast_fp16")]; |
| tensor<fp16, [1, 640]> sigmoid_5_cast_fp16 = sigmoid(x = split_1_cast_fp16_3)[name = string("sigmoid_5_cast_fp16")]; |
| tensor<fp16, [1, 640]> tanh_3_cast_fp16 = tanh(x = add_7_cast_fp16)[name = string("tanh_3_cast_fp16")]; |
| tensor<fp16, [1, 640]> mul_5_cast_fp16 = mul(x = sigmoid_5_cast_fp16, y = tanh_3_cast_fp16)[name = string("mul_5_cast_fp16")]; |
| tensor<fp16, [640, 640]> p_projector_weight_to_fp16 = const()[name = string("p_projector_weight_to_fp16"), val = tensor<fp16, [640, 640]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(23620544)))]; |
| tensor<fp16, [640]> p_projector_bias_to_fp16 = const()[name = string("p_projector_bias_to_fp16"), val = tensor<fp16, [640]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(24439808)))]; |
| tensor<fp16, [1, 640]> pred = linear(bias = p_projector_bias_to_fp16, weight = p_projector_weight_to_fp16, x = mul_5_cast_fp16)[name = string("linear_0_cast_fp16")]; |
| int32 stack_axis_0 = const()[name = string("stack_axis_0"), val = int32(0)]; |
| tensor<fp16, [2, 1, 640]> h_out = stack(axis = stack_axis_0, values = (mul_2_cast_fp16, mul_5_cast_fp16))[name = string("stack_cast_fp16")]; |
| int32 stack_1_axis_0 = const()[name = string("stack_1_axis_0"), val = int32(0)]; |
| tensor<fp16, [2, 1, 640]> c_out = stack(axis = stack_1_axis_0, values = (add_3_cast_fp16, add_7_cast_fp16))[name = string("stack_1_cast_fp16")]; |
| } -> (pred, h_out, c_out); |
| } |
