Repository: tensorflow
Issue: XLA cumsum not supported for int8 and int16
Labels: stat:awaiting tensorflower,type:feature,comp:ops,comp:xla

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

binary

### Tensorflow Version

2.14.0.dev20230503

### Custom Code

No

### OS Platform and Distribution

Google Colab

### Mobile device

_No response_

### Python version

3.10.11

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

`tf.math.cumsum` is not supported for XLA with int8 and int16. I know that not all operations are supported by XLA, but this one seams so trivial. I assume it is an oversight.

### Standalone code to reproduce the issue

Google colab: https://colab.research.google.com/drive/1lqVjSjGmzYwLZfcfgVa8W02lWJ6XcqTk#scrollTo=ex1x-uaCfr6Y

```python
import tensorflow as tf

@tf.function(jit_compile=True)
def cumsum(x):
    return tf.math.cumsum(x)

print(cumsum(tf.constant([0, 1, 2], dtype=tf.dtypes.int8)))
```


### Relevant log output

```shell
InvalidArgumentError: Detected unsupported operations when trying to compile graph __inference_cumsum_7[_XlaMustCompile=true,config_proto=3175580994766145631,executor_type=11160318154034397263] on XLA_CPU_JIT: Cumsum (No registered 'Cumsum' OpKernel for XLA_CPU_JIT devices compatible with node {{node Cumsum}}
	 (OpKernel was found, but attributes didn't match) Requested Attributes: T=DT_INT8, Tidx=DT_INT32, exclusive=false, reverse=false){{node Cumsum}}
The op is created at: 
File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
  return _run_code(code, main_globals, None,
File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
  exec(code, run_globals)
File "/usr/local/lib/python3.10/dist-packages/ipykernel_launcher.py", line 16, in <module>
  app.launch_new_instance()
File "/usr/local/lib/python3.10/dist-packages/traitlets/config/application.py", line 992, in launch_instance
  app.start()
File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelapp.py", line 619, in start
  self.io_loop.start()
File "/usr/local/lib/python3.10/dist-packages/tornado/platform/asyncio.py", line 215, in start
  self.asyncio_loop.run_forever()
File "/usr/lib/python3.10/asyncio/base_events.py", line 603, in run_forever
  self._run_once()
File "/usr/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once
  handle._run()
File "/usr/lib/python3.10/asyncio/events.py", line 80, in _run
  self._context.run(self._callback, *self._args)
File "/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py", line 687, in <lambda>
  lambda f: self._run_callback(functools.partial(callback, future))
File "/usr/local/lib/python3.10/dist-packages/tornado/ioloop.py", line 740, in _run_callback
  ret = callback()
File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 821, in inner
  self.ctx_run(self.run)
File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 782, in run
  yielded = self.gen.send(value)
File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 377, in dispatch_queue
  yield self.process_one()
File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 250, in wrapper
  runner = Runner(ctx_run, result, future, yielded)
File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 748, in __init__
  self.ctx_run(self.run)
File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 782, in run
  yielded = self.gen.send(value)
File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 361, in process_one
  yield gen.maybe_future(dispatch(*args))
File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper
  yielded = ctx_run(next, result)
File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 261, in dispatch_shell
  yield gen.maybe_future(handler(stream, idents, msg))
File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper
  yielded = ctx_run(next, result)
File "/usr/local/lib/python3.10/dist-packages/ipykernel/kernelbase.py", line 539, in execute_request
  self.do_execute(
File "/usr/local/lib/python3.10/dist-packages/tornado/gen.py", line 234, in wrapper
  yielded = ctx_run(next, result)
File "/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py", line 302, in do_execute
  res = shell.run_cell(code, store_history=store_history, silent=silent)
File "/usr/local/lib/python3.10/dist-packages/ipykernel/zmqshell.py", line 539, in run_cell
  return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 2975, in run_cell
  result = self._run_cell(
File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3030, in _run_cell
  return runner(coro)
File "/usr/local/lib/python3.10/dist-packages/IPython/core/async_helpers.py", line 78, in _pseudo_sync_runner
  coro.send(None)
File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3257, in run_cell_async
  has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3473, in run_ast_nodes
  if (await self.run_code(code, result,  async_=asy)):
File "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py", line 3553, in run_code
  exec(code_obj, self.user_global_ns, self.user_ns)
File "<ipython-input-2-602a645c7a37>", line 1, in <cell line: 1>
  print(cumsum(tf.constant([0, 1, 2], dtype=tf.dtypes.int8)))
File "<ipython-input-1-9d904e83e71a>", line 7, in cumsum
  return tf.math.cumsum(x)
	tf2xla conversion failed while converting __inference_cumsum_7[_XlaMustCompile=true,config_proto=3175580994766145631,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_cumsum_7]
```
</details>