From 6b81c0f89864ab4c294a4831facc1cc06cc53544 Mon Sep 17 00:00:00 2001
From: Neil Kichler <neil.kichler@gmail.com>
Date: Fri, 28 Jun 2024 16:39:36 +0200
Subject: [PATCH] Fix type hints in taylor_bounds to support int and float
 arguments.

---
 autobound/jax/jax_bound.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/autobound/jax/jax_bound.py b/autobound/jax/jax_bound.py
index 944648b..eb98249 100644
--- a/autobound/jax/jax_bound.py
+++ b/autobound/jax/jax_bound.py
@@ -58,10 +58,10 @@ def upper(self, x):
 
 
 def taylor_bounds(
-    f: Callable[[jnp.ndarray], jnp.ndarray],
+    f: Callable[[types.NDArrayLike], types.NDArrayLike],
     max_degree: int,
     propagate_trust_regions: bool = False,
-) -> Callable[[jnp.ndarray, tuple[jnp.ndarray, jnp.ndarray]], TaylorBounds]:
+) -> Callable[[types.NDArrayLike, tuple[types.NDArrayLike, types.NDArrayLike]], TaylorBounds]:
   """Returns version of f that returns a TaylorBounds object.
 
   Args: