Repository: tensorflow
Issue: Jit-compiling `tf.while_loop` inside `tf.vectorized_map` raises `InvalidArgumentError`
Labels: stat:awaiting tensorflower,type:bug,comp:apis,comp:ops,2.17

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

binary

### TensorFlow version

v2.17.0

### Custom code

Yes

### OS platform and distribution

Colab

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

MRE
-------
The following mock-up of `cumsum` attempts to JIT compile a `tf.vectorized_map`ped function containing a `tf.scan`:
```python
import tensorflow as tf

def cumsum(xs):
    return tf.scan(
        lambda a, x: a + x, elems=xs
    )

@tf.function(jit_compile=True)
def vec_cumsum(xs):
    return tf.vectorized_map(cumsum, elems=xs)

xs_batched = tf.reshape(tf.range(30), (3, 10))
vec_cumsum(xs_batched)
```

__Expected behaviour__: `vec_cumsum(xs_batched)` returns a batch of cumulative sums.

__Actual behaviour__: Even though all data structures are known statically at JIT compile time, an InvalidArgumentError is raised with "No registered 'TensorListReserve'".  The fault is clearly related to `tf.scan`'s use of `tf.while_loop`, as a (longer) example using naked `tf.while_loop(..., max_iterations=n)` will confirm.

In JAX, it is possible to jit-compile a `vmap`ped function containing a `lax.while_loop` indicating that this is possible in HLO.  It seems the `tf.function(jit_compile=True)` machinery may be mis-transpiling to HLO somehow.

May be related to #73367 also involving `tf.vectorized_map` and `tf.while_loop` (albeit with reversed scope)?

### Standalone code to reproduce the issue

```shell
Colab MRE: https://colab.research.google.com/drive/1bmq1t3PdtebCSlNd0t-iEFrXX7Q0qqZp?usp=sharing
```


### Relevant log output

```shell
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-4-26ea491cd046> in <cell line: 13>()
     11 
     12 # Fails with "No registered 'TensorListReserve'""
---> 13 vec_cumsum(xs_batched)

1 frames
/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     51   try:
     52     ctx.ensure_initialized()
---> 53     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     54                                         inputs, attrs, num_outputs)
     55   except core._NotOkStatusException as e:

InvalidArgumentError: Detected unsupported operations when trying to compile graph __inference_vec_cumsum_462[_XlaMustCompile=true,config_proto=13561319589895757934,executor_type=11160318154034397263] on XLA_CPU_JIT: TensorListReserve (No registered 'TensorListReserve' OpKernel for XLA_CPU_JIT devices compatible with node {{function_node __inference_while_fn_428}}{{node while_init/TensorArrayV2_4}}
	 (OpKernel was found, but attributes didn't match) Requested Attributes: element_dtype=DT_VARIANT, shape_type=DT_INT32){{function_node __inference_while_fn_428}}{{node while_init/TensorArrayV2_4}}
```