Repository: tensorflow
Issue: Gradients of tf.linalg.expm not supported with JIT compilation
Labels: stat:awaiting tensorflower,type:bug,comp:ops,2.17

### 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

```shell
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

```shell
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",
```