From 58b9ec1f6b9c85369349e971cb2e17a016ce35dd Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Wed, 26 Feb 2025 12:16:59 -0800
Subject: [PATCH] kvcache: update tests

---
 kvcache/causal_test.go | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/kvcache/causal_test.go b/kvcache/causal_test.go
index 84d8de54e..412f33e34 100644
--- a/kvcache/causal_test.go
+++ b/kvcache/causal_test.go
@@ -303,6 +303,10 @@ func (b *testBackend) NewContext() ml.Context {
 	return &testContext{}
 }
 
+func (b *testBackend) NewContextSize(int) ml.Context {
+	return &testContext{}
+}
+
 func (b *testBackend) SystemInfo() string {
 	return "not implemented"
 }
@@ -346,11 +350,15 @@ func (c *testContext) FromIntSlice(s []int32, shape ...int) (ml.Tensor, error) {
 	return out, nil
 }
 
+func (c *testContext) Input() ml.Context    { return c }
+func (c *testContext) Output() ml.Context   { return c }
+func (c *testContext) Layer(int) ml.Context { return c }
+
 func (c *testContext) Forward(...ml.Tensor) ml.Context { return c }
 
 func (c *testContext) Compute(...ml.Tensor) {}
 
-func (c *testContext) MaxTensors() int {
+func (c *testContext) MaxGraphNodes() int {
 	return 10
 }