Least squares is 50% slower when activating JIT compile
Hi **@lonidard** ,
I have replicated the reported behavior with jit_compile=True with 2.14,2.13 and tf-nightly. Here I attached a [gist1](https://colab.research.google.com/gist/Venkat6871/1d590ef56e562b957d22fdce5b2a03e6/62378_2-14-2-13-v.ipynb),[ gist2](https://colab.research.google.com/gist/Venkat6871/1d590ef56e562b957d22fdce5b2a03e6/62378_2-14-2-13-v.ipynb) for your reference.

Thank you! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62378">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62378">No</a>
 Thank you Venkat! If there's anything I can do to assist in its resolution, please let me know. Hi @lonidard ,

I have checked the code with Tf2.14v and Tf2.15v and also increased the no of runs and I found in Tf2.15v this API indeed performs better with `jit_compile`. Please find the attached [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/14eb751255b83a3f5c7e6255cfcb0eb5/62378_2-14v-vs-2-15v.ipynb#scrollTo=V54msX3TWU_e) for reference.

```
**In TF2.15v  for 10 runs:**
`lstsq_no_jit`
29.5 ms ± 4.69 ms per loop
`lstsq_jit`
27.6 ms ± 575 µs per loop

**In TF2.14v  for 10 runs:**
`lstsq_no_jit`
20.4 ms ± 2.25 ms per loop
`lstsq_jit`
29.2 ms ± 1.11 ms per loop
```

From the above logs TF2.15v don't have this issue.Thanks! Hi @SuryanarayanaY,

Doesn’t this mean that TF2.15 slowed down the non-jit version without impacting the JIT compiled one, rather than fixing the bug?

For a proper resolution, I’d expect the JIT compiled version to perform better than the no-jit in 2.14. Are there any updates? Hi @lonidard ,

XLA performance depends upon the no of fusions,optimizations for an Op.

Accelerators like GPUs make it even faster.Please refer to attached [gpu-gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/f3e6c8f2d7a5f2d3c181191adc8368af/62378_2-14v-vs-2-15v-r1_gpu.ipynb).

```
**In TF2.15v  for 1000 runs:**
`lstsq_no_jit`
2.29 s ± 346 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
`lstsq_jit`
1.46 s ± 11.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

**In TF2.14v  for 10 runs:**
`lstsq_no_jit`
1.43 s ± 465 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
`lstsq_jit`
1.45 s ± 29.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```

 This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62378">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62378">No</a>

### Issue type

Performance

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

Yes

### Source

source

### TensorFlow version

2.14.0

### Custom code

Yes

### OS platform and distribution

_No response_

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

Activating JIT compile makes tf.linalg.lstsq up to 100% slower. 

If we benchmark `tf.linalg.lstsq` on a matrix of shape (2**18, 5), we see considerable slowdown when activating JIT compile. Expectation would be that JIT actually speeds up solving the least squares problem, but seeing the opposite across multiple platforms and setups. JIT version is consistently slower, taking up to twice as much depending on the platform.

### Standalone code to reproduce the issue


See here:

https://colab.research.google.com/drive/1MIBRIYik2ZVpy1eIIalokMQxE2GA2G9D#scrollTo=7YrOyhUaZD24

or equivalently:

```shell
import tensorflow as tf
@tf.function(jit_compile=False)
def lstsq_no_jit(X,y):
  return tf.linalg.lstsq(X,y)

@tf.function(jit_compile=True)
def lstsq_jit(X,y):
  return tf.linalg.lstsq(X,y)

m, n, k = (2**18, 5, 1)
shape_X = (m, n)
shape_y = (m, k)

X = tf.random.normal(shape_X)
y = tf.random.normal(shape_y)

# Run once to avoid compilation overhead.
lstsq_no_jit(X,y)
lstsq_jit(X,y)
```
At this point, just time
```shell
lstsq_no_jit(X,y)
```
vs 
```shell
lstsq_jit(X,y)
```


### Relevant log output
In the colab version, see the following:
* Running lstsq with no JIT: 6.01 ms ± 401 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
* Running lstsq with JIT: 11.4 ms ± 3.45 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

