diff --git a/code_to_optimize/discrete_riccati.py b/code_to_optimize/discrete_riccati.py
index 53fe30891..b21be8278 100644
--- a/code_to_optimize/discrete_riccati.py
+++ b/code_to_optimize/discrete_riccati.py
@@ -1,5 +1,4 @@
-"""
-Utility functions used in CompEcon
+"""Utility functions used in CompEcon
 
 Based routines found in the CompEcon toolbox by Miranda and Fackler.
 
@@ -9,14 +8,16 @@
 and Finance, MIT Press, 2002.
 
 """
+
 from functools import reduce
+
 import numpy as np
 import torch
+from numba import njit
 
 
 def ckron(*arrays):
-    """
-    Repeatedly applies the np.kron function to an arbitrary number of
+    """Repeatedly applies the np.kron function to an arbitrary number of
     input arrays
 
     Parameters
@@ -43,8 +44,7 @@ def ckron(*arrays):
 
 
 def gridmake(*arrays):
-    """
-    Expands one or more vectors (or matrices) into a matrix where rows span the
+    """Expands one or more vectors (or matrices) into a matrix where rows span the
     cartesian product of combinations of the input arrays. Each column of the
     input arrays will correspond to one column of the output matrix.
 
@@ -79,13 +79,11 @@ def gridmake(*arrays):
                 out = _gridmake2(out, arr)
 
         return out
-    else:
-        raise NotImplementedError("Come back here")
+    raise NotImplementedError("Come back here")
 
 
 def _gridmake2(x1, x2):
-    """
-    Expands two vectors (or matrices) into a matrix where rows span the
+    """Expands two vectors (or matrices) into a matrix where rows span the
     cartesian product of combinations of the input arrays. Each column of the
     input arrays will correspond to one column of the output matrix.
 
@@ -114,19 +112,24 @@ def _gridmake2(x1, x2):
 
     """
     if x1.ndim == 1 and x2.ndim == 1:
-        return np.column_stack([np.tile(x1, x2.shape[0]),
-                               np.repeat(x2, x1.shape[0])])
-    elif x1.ndim > 1 and x2.ndim == 1:
+        # Determine output dtype using numpy's type promotion
+        out_dtype = np.result_type(x1.dtype, x2.dtype)
+        return _gridmake2_optimized_1d_1d(x1, x2, out_dtype)
+    if x1.ndim > 1 and x2.ndim == 1:
+        # For 2D case, use optimized version
+        if x1.ndim == 2:
+            out_dtype = np.result_type(x1.dtype, x2.dtype)
+            return _gridmake2_optimized_2d_1d(x1, x2, out_dtype)
+        # For 3D or higher, fall back to original implementation
+        # This will raise the appropriate error
         first = np.tile(x1, (x2.shape[0], 1))
         second = np.repeat(x2, x1.shape[0])
         return np.column_stack([first, second])
-    else:
-        raise NotImplementedError("Come back here")
+    raise NotImplementedError("Come back here")
 
 
 def _gridmake2_torch(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
-    """
-    PyTorch version of _gridmake2.
+    """PyTorch version of _gridmake2.
 
     Expands two tensors into a matrix where rows span the cartesian product
     of combinations of the input tensors. Each column of the input tensors
@@ -161,10 +164,43 @@ def _gridmake2_torch(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
         first = x1.tile(x2.shape[0])
         second = x2.repeat_interleave(x1.shape[0])
         return torch.column_stack([first, second])
-    elif x1.dim() > 1 and x2.dim() == 1:
+    if x1.dim() > 1 and x2.dim() == 1:
         # tile x1 along first dimension
         first = x1.tile(x2.shape[0], 1)
         second = x2.repeat_interleave(x1.shape[0])
         return torch.column_stack([first, second])
-    else:
-        raise NotImplementedError("Come back here")
+    raise NotImplementedError("Come back here")
+
+
+@njit(cache=True)
+def _gridmake2_optimized_1d_1d(x1: np.ndarray, x2: np.ndarray, out_dtype) -> np.ndarray:
+    """Optimized version for 1D x 1D case."""
+    n1 = x1.shape[0]
+    n2 = x2.shape[0]
+    total = n1 * n2
+
+    out = np.empty((total, 2), dtype=out_dtype)
+    idx = 0
+    for j in range(n2):
+        for i in range(n1):
+            out[idx, 0] = x1[i]
+            out[idx, 1] = x2[j]
+            idx += 1
+    return out
+
+
+@njit(cache=True)
+def _gridmake2_optimized_2d_1d(x1: np.ndarray, x2: np.ndarray, out_dtype) -> np.ndarray:
+    """Optimized version for 2D x 1D case."""
+    rows = x1.shape[0]
+    cols = x1.shape[1]
+    n2 = x2.shape[0]
+    out_rows = rows * n2
+
+    out = np.empty((out_rows, cols + 1), dtype=out_dtype)
+    for j in range(n2):
+        for i in range(rows):
+            for k in range(cols):
+                out[j * rows + i, k] = x1[i, k]
+            out[j * rows + i, cols] = x2[j]
+    return out