From fafb4a5af50a95c8e3c5381f0545b3529358750b Mon Sep 17 00:00:00 2001
From: Pavel Potapov <ppotapov@amazon.com>
Date: Wed, 14 Jan 2026 12:25:29 -0500
Subject: [PATCH] Feat: support __call__ special method

---
 KLR/Trace/NKI.lean | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/KLR/Trace/NKI.lean b/KLR/Trace/NKI.lean
index 8c1a4607..90d0b248 100644
--- a/KLR/Trace/NKI.lean
+++ b/KLR/Trace/NKI.lean
@@ -312,6 +312,13 @@ partial def fnCall
         let args <- bindArgs f (ref :: args) kwargs
         callFn f args
       | _ => throw s!"{func} is not a method of {cls}"
+  | .ref name (.object cls) => do
+      match <- lookup? (.str cls "__call__") with
+      | some (.source f) =>
+        let ref := Term.ref name (.object cls)
+        let args <- bindArgs f (ref :: args) kwargs
+        callFn f args
+      | _ => throw s!"class {cls} is not callable (no __call__ method)"
   | t => throw s!"'{Term.kindStr t}' is not a callable type"
 
 -- Statements