program(1.3)
[IX_checkpoint = string("s3://laticetorch-data/jp/models/ckpt/sydney_lstm_self_attn_m4_laticetorch.pt"), IX_export = string("bolt.apple.com/tasks/brnugxsxez"), IX_zooml = string("https://zooml.apple.com/model/8mnq2zlnq3"), mldb_token = string("mldb-dlqgw2jg9p")]
{
    func batch32<ios17>(tensor<fp32, [1, 15]> history_pad_mask, tensor<fp32, [1, 32]> input, tensor<fp16, [32, 128]> lstm_0_c_in, tensor<fp16, [32, 15, 128]> lstm_0_h_history, tensor<fp16, [32, 128]> lstm_0_h_in) {
            tensor<int32, [1, 32]> input_int32 = cast(dtype = string("int32"), x = input);
            tensor<fp16, [1, 32, 128]> embeddings = gather(axis = int32(0), indices = input_int32, validate_indices = bool(false), x = tensor<fp16, [16032, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(64))));
            tensor<fp16, [1, 32, 128]> lstm_0, tensor<fp16, [32, 128]> lstm_0_h_out, tensor<fp16, [32, 128]> lstm_0_c_out = lstm(activation = string("tanh"), bias = tensor<fp16, [512]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4104320))), cell_activation = string("tanh"), direction = string("forward"), initial_c = lstm_0_c_in, initial_h = lstm_0_h_in, output_sequence = bool(true), recurrent_activation = string("sigmoid"), weight_hh = tensor<fp16, [512, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4105408))), weight_ih = tensor<fp16, [512, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4236544))), x = embeddings);
            tensor<fp16, [1, 15]> history_pad_mask_fp16 = cast(dtype = string("fp16"), x = history_pad_mask);
            tensor<fp16, [1, 1, 15]> history_pad_mask_3d = expand_dims(axes = tensor<int32, [1]>([1]), x = history_pad_mask_fp16);
            tensor<fp16, [1, 1, 15]> history_softmax_mask = mul(x = history_pad_mask_3d, y = fp16(-10000));
            tensor<fp16, [32, 1, 128]> lstm_0_T = transpose(perm = tensor<int32, [3]>([1, 0, 2]), x = lstm_0);
            tensor<fp16, [32, 1, 15]> self_attn_0_dot = matmul(transpose_x = bool(false), transpose_y = bool(true), x = lstm_0_T, y = lstm_0_h_history);
            tensor<fp16, [32, 1, 15]> self_attn_0_masked = add(x = self_attn_0_dot, y = history_softmax_mask);
            tensor<fp16, [32, 1, 15]> self_attn_0_weights = softmax(axis = int32(-1), x = self_attn_0_masked);
            tensor<fp16, [32, 1, 128]> self_attn_0_ctx_T = matmul(transpose_x = bool(false), transpose_y = bool(false), x = self_attn_0_weights, y = lstm_0_h_history);
            tensor<fp16, [1, 32, 128]> self_attn_0_ctx = transpose(perm = tensor<int32, [3]>([1, 0, 2]), x = self_attn_0_ctx_T);
            tensor<fp16, [1, 32, 256]> self_attn_0_h_ctx = concat(axis = int32(-1), interleave = bool(false), values = (lstm_0, self_attn_0_ctx));
            tensor<fp16, [1, 32, 128]> self_attn_0_proj = linear(weight = tensor<fp16, [128, 256]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4367680))), x = self_attn_0_h_ctx);
            tensor<fp16, [1, 32, 128]> self_attn_0 = tanh(x = self_attn_0_proj);
            tensor<fp16, [1, 32, 16032]> logits = linear(bias = tensor<fp16, [16032]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4433280))), weight = tensor<fp16, [16032, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4465408))), x = self_attn_0);
        } -> (lstm_0_h_out, lstm_0_c_out, logits);
    func logJointProb<ios17>(tensor<fp32, [1, 15]> history_pad_mask, tensor<fp32, [1, 1]> input, tensor<fp32, [1, 1]> label, tensor<fp16, [1, 128]> lstm_0_c_in, tensor<fp16, [1, 15, 128]> lstm_0_h_history, tensor<fp16, [1, 128]> lstm_0_h_in) {
            tensor<int32, [1, 1]> input_int32 = cast(dtype = string("int32"), x = input);
            tensor<fp16, [1, 1, 128]> embeddings = gather(axis = int32(0), indices = input_int32, validate_indices = bool(false), x = tensor<fp16, [16032, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(64))));
            tensor<fp16, [1, 1, 128]> lstm_0, tensor<fp16, [1, 128]> lstm_0_h_out, tensor<fp16, [1, 128]> lstm_0_c_out = lstm(activation = string("tanh"), bias = tensor<fp16, [512]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4104320))), cell_activation = string("tanh"), direction = string("forward"), initial_c = lstm_0_c_in, initial_h = lstm_0_h_in, output_sequence = bool(true), recurrent_activation = string("sigmoid"), weight_hh = tensor<fp16, [512, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4105408))), weight_ih = tensor<fp16, [512, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4236544))), x = embeddings);
            tensor<fp16, [1, 15]> history_pad_mask_fp16 = cast(dtype = string("fp16"), x = history_pad_mask);
            tensor<fp16, [1, 1, 15]> history_pad_mask_3d = expand_dims(axes = tensor<int32, [1]>([1]), x = history_pad_mask_fp16);
            tensor<fp16, [1, 1, 15]> history_softmax_mask = mul(x = history_pad_mask_3d, y = fp16(-10000));
            tensor<fp16, [1, 1, 128]> lstm_0_T = transpose(perm = tensor<int32, [3]>([1, 0, 2]), x = lstm_0);
            tensor<fp16, [1, 1, 15]> self_attn_0_dot = matmul(transpose_x = bool(false), transpose_y = bool(true), x = lstm_0_T, y = lstm_0_h_history);
            tensor<fp16, [1, 1, 15]> self_attn_0_masked = add(x = self_attn_0_dot, y = history_softmax_mask);
            tensor<fp16, [1, 1, 15]> self_attn_0_weights = softmax(axis = int32(-1), x = self_attn_0_masked);
            tensor<fp16, [1, 1, 128]> self_attn_0_ctx_T = matmul(transpose_x = bool(false), transpose_y = bool(false), x = self_attn_0_weights, y = lstm_0_h_history);
            tensor<fp16, [1, 1, 128]> self_attn_0_ctx = transpose(perm = tensor<int32, [3]>([1, 0, 2]), x = self_attn_0_ctx_T);
            tensor<fp16, [1, 1, 256]> self_attn_0_h_ctx = concat(axis = int32(-1), interleave = bool(false), values = (lstm_0, self_attn_0_ctx));
            tensor<fp16, [1, 1, 128]> self_attn_0_proj = linear(weight = tensor<fp16, [128, 256]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4367680))), x = self_attn_0_h_ctx);
            tensor<fp16, [1, 1, 128]> self_attn_0 = tanh(x = self_attn_0_proj);
            tensor<int32, [1, 1]> label_int32 = cast(dtype = string("int32"), x = label);
            tensor<fp16, [1, 1, 128]> label_embeddings = gather(axis = int32(0), indices = label_int32, validate_indices = bool(false), x = tensor<fp16, [16032, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4465408))));
            tensor<fp16, [1, 1, 128, 1]> label_embeddings_nd = expand_dims(axes = tensor<int32, [1]>([-1]), x = label_embeddings);
            tensor<fp16, [1, 1, 128, 1]> model_embeddings_nd = expand_dims(axes = tensor<int32, [1]>([-1]), x = self_attn_0);
            tensor<fp16, [1, 1, 1, 1]> logits_without_biases_nd = matmul(transpose_x = bool(true), transpose_y = bool(false), x = model_embeddings_nd, y = label_embeddings_nd);
            tensor<fp16, [1, 1]> logits_without_biases = squeeze(axes = tensor<int32, [2]>([-1, -2]), x = logits_without_biases_nd);
            tensor<fp16, [1, 1]> logits_biases = gather(axis = int32(0), indices = label_int32, validate_indices = bool(false), x = tensor<fp16, [16032]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4433280))));
            tensor<fp16, [1, 1]> logits = add(x = logits_without_biases, y = logits_biases);
        } -> (lstm_0_h_out, lstm_0_c_out, logits);
    func logJointProb_batch32<ios17>(tensor<fp32, [1, 15]> history_pad_mask, tensor<fp32, [1, 32]> input, tensor<fp32, [1, 32]> label, tensor<fp16, [32, 128]> lstm_0_c_in, tensor<fp16, [32, 15, 128]> lstm_0_h_history, tensor<fp16, [32, 128]> lstm_0_h_in) {
            tensor<int32, [1, 32]> input_int32 = cast(dtype = string("int32"), x = input);
            tensor<fp16, [1, 32, 128]> embeddings = gather(axis = int32(0), indices = input_int32, validate_indices = bool(false), x = tensor<fp16, [16032, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(64))));
            tensor<fp16, [1, 32, 128]> lstm_0, tensor<fp16, [32, 128]> lstm_0_h_out, tensor<fp16, [32, 128]> lstm_0_c_out = lstm(activation = string("tanh"), bias = tensor<fp16, [512]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4104320))), cell_activation = string("tanh"), direction = string("forward"), initial_c = lstm_0_c_in, initial_h = lstm_0_h_in, output_sequence = bool(true), recurrent_activation = string("sigmoid"), weight_hh = tensor<fp16, [512, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4105408))), weight_ih = tensor<fp16, [512, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4236544))), x = embeddings);
            tensor<fp16, [1, 15]> history_pad_mask_fp16 = cast(dtype = string("fp16"), x = history_pad_mask);
            tensor<fp16, [1, 1, 15]> history_pad_mask_3d = expand_dims(axes = tensor<int32, [1]>([1]), x = history_pad_mask_fp16);
            tensor<fp16, [1, 1, 15]> history_softmax_mask = mul(x = history_pad_mask_3d, y = fp16(-10000));
            tensor<fp16, [32, 1, 128]> lstm_0_T = transpose(perm = tensor<int32, [3]>([1, 0, 2]), x = lstm_0);
            tensor<fp16, [32, 1, 15]> self_attn_0_dot = matmul(transpose_x = bool(false), transpose_y = bool(true), x = lstm_0_T, y = lstm_0_h_history);
            tensor<fp16, [32, 1, 15]> self_attn_0_masked = add(x = self_attn_0_dot, y = history_softmax_mask);
            tensor<fp16, [32, 1, 15]> self_attn_0_weights = softmax(axis = int32(-1), x = self_attn_0_masked);
            tensor<fp16, [32, 1, 128]> self_attn_0_ctx_T = matmul(transpose_x = bool(false), transpose_y = bool(false), x = self_attn_0_weights, y = lstm_0_h_history);
            tensor<fp16, [1, 32, 128]> self_attn_0_ctx = transpose(perm = tensor<int32, [3]>([1, 0, 2]), x = self_attn_0_ctx_T);
            tensor<fp16, [1, 32, 256]> self_attn_0_h_ctx = concat(axis = int32(-1), interleave = bool(false), values = (lstm_0, self_attn_0_ctx));
            tensor<fp16, [1, 32, 128]> self_attn_0_proj = linear(weight = tensor<fp16, [128, 256]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4367680))), x = self_attn_0_h_ctx);
            tensor<fp16, [1, 32, 128]> self_attn_0 = tanh(x = self_attn_0_proj);
            tensor<int32, [1, 32]> label_int32 = cast(dtype = string("int32"), x = label);
            tensor<fp16, [1, 32, 128]> label_embeddings = gather(axis = int32(0), indices = label_int32, validate_indices = bool(false), x = tensor<fp16, [16032, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4465408))));
            tensor<fp16, [1, 32, 128, 1]> label_embeddings_nd = expand_dims(axes = tensor<int32, [1]>([-1]), x = label_embeddings);
            tensor<fp16, [1, 32, 128, 1]> model_embeddings_nd = expand_dims(axes = tensor<int32, [1]>([-1]), x = self_attn_0);
            tensor<fp16, [1, 32, 1, 1]> logits_without_biases_nd = matmul(transpose_x = bool(true), transpose_y = bool(false), x = model_embeddings_nd, y = label_embeddings_nd);
            tensor<fp16, [1, 32]> logits_without_biases = squeeze(axes = tensor<int32, [2]>([-1, -2]), x = logits_without_biases_nd);
            tensor<fp16, [1, 32]> logits_biases = gather(axis = int32(0), indices = label_int32, validate_indices = bool(false), x = tensor<fp16, [16032]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4433280))));
            tensor<fp16, [1, 32]> logits = add(x = logits_without_biases, y = logits_biases);
        } -> (lstm_0_h_out, lstm_0_c_out, logits);
    func main<ios17>(tensor<fp32, [1, 15]> history_pad_mask, tensor<fp32, [1, 1]> input, tensor<fp16, [1, 128]> lstm_0_c_in, tensor<fp16, [1, 15, 128]> lstm_0_h_history, tensor<fp16, [1, 128]> lstm_0_h_in) {
            tensor<int32, [1, 1]> input_int32 = cast(dtype = string("int32"), x = input);
            tensor<fp16, [1, 1, 128]> embeddings = gather(axis = int32(0), indices = input_int32, validate_indices = bool(false), x = tensor<fp16, [16032, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(64))));
            tensor<fp16, [1, 1, 128]> lstm_0, tensor<fp16, [1, 128]> lstm_0_h_out, tensor<fp16, [1, 128]> lstm_0_c_out = lstm(activation = string("tanh"), bias = tensor<fp16, [512]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4104320))), cell_activation = string("tanh"), direction = string("forward"), initial_c = lstm_0_c_in, initial_h = lstm_0_h_in, output_sequence = bool(true), recurrent_activation = string("sigmoid"), weight_hh = tensor<fp16, [512, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4105408))), weight_ih = tensor<fp16, [512, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4236544))), x = embeddings);
            tensor<fp16, [1, 15]> history_pad_mask_fp16 = cast(dtype = string("fp16"), x = history_pad_mask);
            tensor<fp16, [1, 1, 15]> history_pad_mask_3d = expand_dims(axes = tensor<int32, [1]>([1]), x = history_pad_mask_fp16);
            tensor<fp16, [1, 1, 15]> history_softmax_mask = mul(x = history_pad_mask_3d, y = fp16(-10000));
            tensor<fp16, [1, 1, 128]> lstm_0_T = transpose(perm = tensor<int32, [3]>([1, 0, 2]), x = lstm_0);
            tensor<fp16, [1, 1, 15]> self_attn_0_dot = matmul(transpose_x = bool(false), transpose_y = bool(true), x = lstm_0_T, y = lstm_0_h_history);
            tensor<fp16, [1, 1, 15]> self_attn_0_masked = add(x = self_attn_0_dot, y = history_softmax_mask);
            tensor<fp16, [1, 1, 15]> self_attn_0_weights = softmax(axis = int32(-1), x = self_attn_0_masked);
            tensor<fp16, [1, 1, 128]> self_attn_0_ctx_T = matmul(transpose_x = bool(false), transpose_y = bool(false), x = self_attn_0_weights, y = lstm_0_h_history);
            tensor<fp16, [1, 1, 128]> self_attn_0_ctx = transpose(perm = tensor<int32, [3]>([1, 0, 2]), x = self_attn_0_ctx_T);
            tensor<fp16, [1, 1, 256]> self_attn_0_h_ctx = concat(axis = int32(-1), interleave = bool(false), values = (lstm_0, self_attn_0_ctx));
            tensor<fp16, [1, 1, 128]> self_attn_0_proj = linear(weight = tensor<fp16, [128, 256]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4367680))), x = self_attn_0_h_ctx);
            tensor<fp16, [1, 1, 128]> self_attn_0 = tanh(x = self_attn_0_proj);
            tensor<fp16, [1, 1, 16032]> logits = linear(bias = tensor<fp16, [16032]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4433280))), weight = tensor<fp16, [16032, 128]>(BLOBFILE(path = string("@model_path/lstm.mil.weights"), offset = uint64(4465408))), x = self_attn_0);
        } -> (lstm_0_h_out, lstm_0_c_out, logits);
}