compiled/moss_token_embedding.mlmodelc/model.mil

program(1.0)
[buildInfo = dict<tensor<string, []>, tensor<string, []>>({{"coremlc-component-MIL", "3520.4.1"}, {"coremlc-version", "3520.5.1"}, {"coremltools-component-torch", "2.12.1"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "9.0"}})]
{
    func main<ios17>(tensor<int32, [1, 512]> input_ids) {
            tensor<fp16, [151936, 2048]> embedding_weight = const()[name = tensor<string, []>("embedding_weight"), val = tensor<fp16, [151936, 2048]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(64)))];
            tensor<int32, []> var_11_batch_dims_0 = const()[name = tensor<string, []>("op_11_batch_dims_0"), val = tensor<int32, []>(0)];
            tensor<bool, []> var_11_validate_indices_0 = const()[name = tensor<string, []>("op_11_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, 512]> greater_equal_0 = greater_equal(x = input_ids, 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, []>(151936)];
            tensor<int32, [1, 512]> add_0 = add(x = input_ids, y = slice_by_index_0)[name = tensor<string, []>("add_0")];
            tensor<int32, [1, 512]> select_0 = select(a = input_ids, b = add_0, cond = greater_equal_0)[name = tensor<string, []>("select_0")];
            tensor<int32, []> var_11_axis_1 = const()[name = tensor<string, []>("op_11_axis_1"), val = tensor<int32, []>(0)];
            tensor<fp16, [1, 512, 2048]> token_embeddings = gather(axis = var_11_axis_1, batch_dims = var_11_batch_dims_0, indices = select_0, validate_indices = var_11_validate_indices_0, x = embedding_weight)[name = tensor<string, []>("op_11")];
        } -> (token_embeddings);
}