Check if the Triton dialect bindings are available in lib/triton.py
IIRC we used to import these bindings in lib/__init__.py which is imported
as part of the top-level jax package. So, it did make sense to delay the
check until we actually need the bindings.
However, we have since moved the bindings to lib/triton.py and thus we could
move the check there.
PiperOrigin-RevId: 607196039
