When cuda is not available, the attributes `triton_kernel_call_lib.Tr…

…itonKernel` and `triton.JITFunction` are missing. This creates issues when trying to import the code, e.g. using colab adhoc impoort.

Delaying the evaluation of these types fixes the problem.

PiperOrigin-RevId: 568330491

jax_triton/triton_lib.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 
 """Module for calling Triton kernels from JAX."""
+# b/301982023
+from __future__ import annotations
+
 import functools
 import os
 import types