From 16e066ec692f0a2d3cb31a0638a02ba89cb72687 Mon Sep 17 00:00:00 2001
From: Foke Singh <foke.singh@deepkapha.com>
Date: Wed, 11 Jan 2023 09:58:09 +0000
Subject: [PATCH] Add llama example demonstrating the new gatherValues
 functionality.

---
 examples/llama/llama.zig | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/llama/llama.zig b/examples/llama/llama.zig
index 0112c01..f01ab91 100644
--- a/examples/llama/llama.zig
+++ b/examples/llama/llama.zig
@@ -77,7 +77,7 @@ pub const LlamaLM = struct {
         const tokens = tokens_.withPartialTags(.{.s});
         const out = out_.withPartialTags(.{ .s, .d });
 
-        const next_token_pred = out.dynamicSlice(.{ .s = .{ .start = token_index, .len = 1 } });
+        const next_token_pred = out.gatherValues(.s, token_index, .{});
         var logits = zml.call(lm_head, .forward, .{next_token_pred});
         if (logits.shape().hasTag(.voc) == null)
             logits = logits.rename(.{ .d = .voc });