Fix missing tree_reduce import in models/cache.py (#1165)

2026-04-20 13:31:58 -05:00
parent 6d114686e5
commit a401730941
1 changed files with 1 additions and 1 deletions
@@ -7,7 +7,7 @@ from typing import Any, Dict, List, Optional

 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_map, tree_unflatten
+from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten

 from .base import create_causal_mask