Gradients of tf.linalg.expm not supported with JIT compilation

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

binary

TensorFlow version

tested 2.17 and 2.10, both have the issue

Custom code

Yes

OS platform and distribution

Ubuntu

Mobile device

No response

Python version

tested 3.9 and 3.12

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

Gradients of tf.linalg.expm can not be computed with JIT compilation.

This is an issue, because tf 2.17 seems to have activated jit compilation for compiled models per default whereas earlier versions did not, breaking existing code.

Standalone code to reproduce the issue

import tensorflow as tf

A = tf.Variable([[.4, 1.5], [.6, .1]], dtype=tf.float32)

@tf.function(jit_compile=True) #set jit_compile=False to make it work
def f(A):
    with tf.GradientTape() as tape:
        B = tf.linalg.expm(A)
    return tape.gradient(B, A)

f(A)

Relevant log output

2024-10-11 11:17:27.281304: W tensorflow/core/framework/op_kernel.cc:1840] OP_REQUIRES failed at xla_ops.cc:577 : INVALID_ARGUMENT: XLA compilation requires a fixed tensor list size. Set the max number of elements. This could also happen if you're using a TensorArray in a while loop that does not have its maximum_iteration set, you can fix this by setting maximum_iteration to a suitable value.

Stack trace for op definition: 
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/site-packages/ipykernel_launcher.py", line 18, in <module>
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/site-packages/traitlets/config/application.py", line 1075, in launch_instance
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/site-packages/ipykernel/kernelapp.py", line 739, in start
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/site-packages/tornado/platform/asyncio.py", line 205, in start
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/asyncio/base_events.py", line 641, in run_forever
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/asyncio/base_events.py", line 1986, in _run_once
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/asyncio/events.py", line 88, in _run
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 545, in dispatch_queue
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 534, in process_one
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/site-packages/ipykernel/kernelbase.py", line 437, in dispatch_shell
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/site-packages/ipykernel/ipkernel.py", line 362, in execute_request
File "/home/beckerf/mambaforge/envs/learnMSAdev2/lib/python3.12/site-packages/ipykernel/kernelbase.py",