Fix type hints in taylor_bounds to support int and float arguments.

[neilkichler]

In the README example, Pyright complains about the arguments to taylor_bounds since the types of the Callable are restricted to jnp.ndarray and not types.NDArraylike which also supports int, float etc.