Repository: tensorflow
Issue: tf.sparse.reduce_sum error in JIT
Labels: stat:awaiting tensorflower,type:bug,comp:ops,comp:xla,2.17

### Issue type

Bug

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

No

### Source

binary

### TensorFlow version

v2.17.0-rc1-2-gad6d8cc177d

### Custom code

Yes

### OS platform and distribution

Ubuntu Mate 22.04

### Mobile device

_No response_

### Python version

Python 3.10.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

CUDA 12.3

### GPU model and memory

_No response_

### Current behavior?

Error when using tf.sparse.reduce_sum and JIT compilation.

I have written a layer that passes all my unit tests except when I use in a model with predict or train. 
Would that make sense with the JIT compilation on? I am not fully certain how and when this works.

Anyway, I am not exactly sure why my layer code fails, but I think that this minimum reproducible example captures the issue.

### Standalone code to reproduce the issue

```shell
import numpy as np
import tensorflow as tf

batch_size = 4
input_shape = (3, 3)

indices = np.array([[0, 0, 0], [0, 0, 1], [1, 0, 0], [1, 0, 1], [2, 0, 0], [2, 0, 1], [3, 0, 0], [3, 0, 1]])
inputs = tf.sparse.SparseTensor(dense_shape=(batch_size, *input_shape),
                                indices=indices,
                                values=[9, 1, 9, 1, 9, 1, 9, 1])


@tf.function(input_signature=[tf.SparseTensorSpec(
    shape=(4, 3, 3),
    dtype=tf.dtypes.int32)], jit_compile=True)
def get_batch_sum(inputs):
    # same problem with tf.sparse.reduce_max
    return tf.sparse.reduce_sum(tf.sparse.reduce_sum(inputs, axis=1, output_is_sparse=True), axis=1)


sum_out = get_batch_sum(inputs)
print(sum_out)
```


### Relevant log output

```shell
2024-08-20 16:31:01.690779: W tensorflow/core/framework/op_kernel.cc:1840] OP_REQUIRES failed at xla_ops.cc:577 : INVALID_ARGUMENT: Detected unsupported operations when trying to compile graph __inference_get_batch_sum_15[_XlaMustCompile=true,config_proto=6001324581131673121,executor_type=11160318154034397263] on XLA_GPU_JIT: SparseReduceSumSparse (No registered 'SparseReduceSumSparse' OpKernel for XLA_GPU_JIT devices compatible with node {{node SparseReduceSumSparse}}){{node SparseReduceSumSparse}}
The op is created at: 
File ".config/JetBrains/PyCharmCE2023.3/scratches/scratch_30.py", line 21, in <module>
File ".config/JetBrains/PyCharmCE2023.3/scratches/scratch_30.py", line 18, in get_batch_sum
	tf2xla conversion failed while converting __inference_get_batch_sum_15[_XlaMustCompile=true,config_proto=6001324581131673121,executor_type=11160318154034397263]. Run with TF_DUMP_GRAPH_PREFIX=/path/to/dump/dir and --vmodule=xla_compiler=2 to obtain a dump of the compiled functions.
Traceback (most recent call last):
(...)  

File ".config/JetBrains/PyCharmCE2023.3/scratches/scratch_30.py", line 18, in get_batch_sum
	tf2xla conversion failed while converting __inference_get_batch_sum_15[_XlaMustCompile=true,config_proto=6001324581131673121,executor_type=11160318154034397263]. Run with TF_DUMP_GRAPH_PREFIX=/path/to/dump/dir and --vmodule=xla_compiler=2 to obtain a dump of the compiled functions. [Op:__inference_get_batch_sum_15]
```