InternalError: libdevice not found at ./libdevice.10.bc

Hello,
I have some problems when I try to implement jit compilation for an ode. The minimum code to reproduce is

@tf.function(jit_compile=True)
def tf_spre(O):
    A = tf.eye( np.prod(O.shape[1]), dtype = tf.complex128 )
    B = O
    N = A.shape[0]
    return tf.reshape(tf.einsum('ik,jl', A, B), [N**2, N**2])

@tf.function(jit_compile=True)
def tf_spost(O):
    A = tf.transpose(O)
    B = tf.eye( np.prod(O.shape[0]), dtype = tf.complex128 )
    N = A.shape[0]
    return tf.reshape(tf.einsum('ik,jl', A, B), [N**2, N**2])

@tf.function(jit_compile=True)
def tf_sprepost(A, B):
    return tf.tensordot(tf_spre(A), tf_spost(B), axes = 1)

@tf.function(jit_compile=True)
def lindblad_dissipator(c_op):
    prod = tf.tensordot( tfla.adjoint(c_op), c_op, axes = 1 )
    return tf_sprepost(c_op, tfla.adjoint(c_op)) - 0.5 * tf_spre(prod) - 0.5 * tf_spost(prod)

@tf.function(jit_compile=True)
def tf_liouvillian(H):
    L = -1j * (tf_spre(H) - tf_spost(H))
    return L

@tf.function(jit_compile=True)
def expval(Op, rho_s):
    tmp = tf.tensordot( rho_s, Op, axes = [[2], [0]])
    return tfmt.real( tf.einsum('ijj->i', tmp) )

@tf.function(jit_compile=True)
def drho_dt(t, rho, H, H_t, lindbladian):
    A_sin = 0.01 * tfmt.sin(t)
    H_t = H_t * tf.cast(A_sin, dtype = tf.complex128)
    return tf.tensordot(tf_liouvillian(H + H_t) + lindbladian, rho, axes = [[1], [0]])

def tf_mesolve(H, H_t, psi0, t_l, c_ops, e_ops):
    H = tf.constant( H.full(), dtype = tf.complex128 )
    H_t = tf.constant( H_t.full(), dtype = tf.complex128 )
    psi0 = tf.constant( psi0.full(), dtype = tf.complex128 )
    rho0 = tf.tensordot(psi0, tfla.adjoint(psi0), axes = 1)
    
    N = H.shape[0]
    
    if len(c_ops) != 0:
        c_ops2 = np.zeros([len(c_ops), N, N], dtype = np.complex128)
        for i in range(len(c_ops)):
            c_ops2[i] = np.array( c_ops[i].full() )
    else:
        c_ops2 = np.zeros([1, N, N], dtype = np.complex128)
    
    if len(e_ops) != 0:
        e_ops2 = np.zeros([len(e_ops), N, N], dtype = np.complex128)
        for i in range(len(e_ops)):
            e_ops2[i] = np.array( e_ops[i].full() )
    else:
        e_ops2 = np.zeros([1, N, N], dtype = np.complex128)
    
    lindbladian = np.zeros([N**2, N**2], dtype = np.complex128)
    for i in range(c_ops2.shape[0]):
        lindbladian += lindblad_dissipator(c_ops2[i])
        
    #run_fn = tf.function(lambda: tfp.math.ode.DormandPrince(rtol=1e-3, atol=1e-5).solve(drho_dt, t_l[0], tf.reshape(rho0, [N**2]),
    #                               solution_times=t_l,
    #                                constants={"H": H, "H_t": H_t, "lindbladian": lindbladian}), jit_compile=True)
    #results = run_fn()
    
    results = tfp.math.ode.DormandPrince(rtol=1e-3, atol=1e-5, first_step_size=1e-2, safety_factor=0.9,
    min_step_size_factor=0.1, max_step_size_factor=10.0).solve(drho_dt, t_l[0], tf.reshape(rho0, [N**2]),
                                   solution_times=tfp.math.ode.ChosenBySolver(final_time=t_l[-1]), constants={"H": H, "H_t": H_t, "lindbladian": lindbladian})
    
    
    times = results.times
    states = results.states
    
    states = tf.reshape(states, [times.shape[0], N, N])
    
    expvals = []
    for i in range(len(e_ops)):
        expvals.append( expval(e_ops2[i], states) )
        
    return times, states, expvals

w_q = 1

sx = qtp.sigmax()
sz = qtp.sigmaz()
sp = qtp.sigmap()
sm = qtp.sigmam()

H = 0.5 * w_q * sz

gam_q = 0.1

c_ops = []
c_ops.append(np.sqrt(gam_q) * sm)
#c_ops.append(np.sqrt(gam_q) * sz)

e_ops = []
e_ops.append(sp * sm)

psi0 = (qtp.fock(2, 0) + 1j * qtp.fock(2, 1)).unit()

t_l = np.linspace(0, 10 / gam_q, 100)

times, states, tf_expect = tf_mesolve(H, sm + sp, psi0, t_l, c_ops, e_ops)

Here H is a 2x2 matrix.
This code returns the following error:

InternalError                             Traceback (most recent call last)
<timed exec> in <module>

<ipython-input-3-897b1854b462> in tf_mesolve(H, H_t, psi0, t_l, c_ops, e_ops)
     69     #results = run_fn()
     70 
---> 71     results = tfp.math.ode.DormandPrince(rtol=1e-3, atol=1e-5, first_step_size=1e-2, safety_factor=0.9,
     72     min_step_size_factor=0.1, max_step_size_factor=10.0).solve(drho_dt, t_l[0], tf.reshape(rho0, [N**2]),
     73                                    solution_times=tfp.math.ode.ChosenBySolver(final_time=t_l[-1]), constants={"H": H, "H_t": H_t, "lindbladian": lindbladian})

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow_probability/python/math/ode/base.py in solve(self, ode_fn, initial_time, initial_state, solution_times, jacobian_fn, jacobian_sparsity, batch_ndims, previous_solver_internal_state, constants)
    473     # custom_gradient will complain even if there are no variables in `ode_fn`.
    474     with tf1.variable_scope(tf1.get_variable_scope(), use_resource=True):
--> 475       return gradient_helper(*(flat_initial_state + flat_constants))
    476 
    477   @abc.abstractmethod

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow/python/ops/custom_gradient.py in __call__(self, *a, **k)
    259 
    260   def __call__(self, *a, **k):
--> 261     return self._d(self._f, a, k)
    262 
    263 

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow/python/ops/custom_gradient.py in decorated(wrapped, args, kwargs)
    213 
    214     if context.executing_eagerly():
--> 215       return _eager_mode_decorator(wrapped, args, kwargs)
    216     else:
    217       return _graph_mode_decorator(wrapped, args, kwargs)

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow/python/ops/custom_gradient.py in _eager_mode_decorator(f, args, kwargs)
    437   """Implement custom gradient decorator for eager mode."""
    438   with tape_lib.VariableWatcher() as variable_watcher:
--> 439     result, grad_fn = f(*args, **kwargs)
    440   args = nest.flatten(args)
    441   all_inputs = list(args) + list(kwargs.values())

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow_probability/python/math/ode/base.py in gradient_helper(*flat_initial_state_and_constants)
    229           constant_state_structure, flat_constants)
    230 
--> 231       results = self._solve(
    232           ode_fn=functools.partial(ode_fn, **constants),
    233           initial_time=initial_time,

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow_probability/python/math/ode/dormand_prince.py in _solve(***failed resolving arguments***)
    197       solver_internal_state = previous_solver_internal_state
    198       if solver_internal_state is None:
--> 199         solver_internal_state = self._initialize_solver_internal_state(
    200             ode_fn=ode_fn,
    201             initial_state=p.initial_state,

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow_probability/python/math/ode/dormand_prince.py in _initialize_solver_internal_state(self, ode_fn, initial_time, initial_state)
    307     p = self._prepare_common_params(initial_state, initial_time)
    308 
--> 309     initial_derivative = ode_fn(p.initial_time, p.initial_state)
    310     initial_derivative = tf.nest.map_structure(tf.convert_to_tensor,
    311                                                initial_derivative)

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    887 
    888       with OptionalXlaContext(self._jit_compile):
--> 889         result = self._call(*args, **kwds)
    890 
    891       new_tracing_count = self.experimental_get_tracing_count()

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    954               *args, **kwds)
    955       # If we did not create any variables the trace we have is good enough.
--> 956       return self._concrete_stateful_fn._call_flat(
    957           filtered_flat_args, self._concrete_stateful_fn.captured_inputs)  # pylint: disable=protected-access
    958 

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow/python/eager/function.py in _call_flat(self, args, captured_inputs, cancellation_manager)
   1958         and executing_eagerly):
   1959       # No tape is watching; skip to running the function.
-> 1960       return self._build_call_outputs(self._inference_function.call(
   1961           ctx, args, cancellation_manager=cancellation_manager))
   1962     forward_backward = self._select_forward_and_backward_functions(

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow/python/eager/function.py in call(self, ctx, args, cancellation_manager)
    589       with _InterpolateFunctionError(self):
    590         if cancellation_manager is None:
--> 591           outputs = execute.execute(
    592               str(self.signature.name),
    593               num_outputs=self._num_outputs,

~/anaconda3/envs/physics/lib/python3.9/site-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     57   try:
     58     ctx.ensure_initialized()
---> 59     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     60                                         inputs, attrs, num_outputs)
     61   except core._NotOkStatusException as e:

InternalError: libdevice not found at ./libdevice.10.bc [Op:__inference_drho_dt_12917]

I noticed it comes from a conflict between @tf.function(jit_compile=True) method and the time variable inside the drho_dt function. Indeed, if i eliminate the cast of the explicit time dependence, that is

@tf.function(jit_compile=True)
def drho_dt(t, rho, H, H_t, lindbladian):
    A_sin = 0.01 * tfmt.sin(t)
    return tf.tensordot(tf_liouvillian(H + H_t) + lindbladian, rho, axes = [[1], [0]])

of if I eliminate the @tf.function(jit_compile=True) method, the problem vanishes. Of course i need the time dependence for simulate driven dissipative quantum systems, and i want to compile that function, to speed up the calculation.

Moreover, the tensorflow ode is faster than the ode implemented in qutip package only for this small system. If I enter for example 20x20 matrices, the qutip (that is cpu) ode is 10 times faster.

Start with the minimal Python reproduction around drho_dt and tfp.math.ode.DormandPrince, comparing the jit_compile=True and non-JIT paths. Reproduce the libdevice.10.bc error and determine whether the issue is in TensorFlow Probability's ODE invocation or the TensorFlow/XLA environment; done means the time-dependent JIT path works or the required environment limitation is documented.