C++ API `MatrixDiagV3` aborts with tensor type of `k`
@Sehun0819 If MatrixDiagV3 doesn't work due to type limitations, could you consider alternative approaches to achieve the desired functionality like tf.linalg.diag?
Thank you! @sushreebarsa 
Hi! Even there exist alternatives this should be patched to avoid crash right?
Please check python version of [reproduction](https://colab.research.google.com/drive/1R11FohcZZ-Ky1rpO-bJPOcztHdRIQE5R?usp=sharing).
```Python
import tensorflow as tf

tf.raw_ops.MatrixDiagV3(
    diagonal=tf.constant([0.0],dtype=tf.float32),
    k=tf.constant([1,2],shape=[2],dtype=tf.int32),
    num_rows=1,
    num_cols=1,
    padding_value=0.0,
    align='RIGHT_LEFT',
    name=None
)
``` Hi @Sehun0819 ,

Thanks for reporting. Here the checkfail is happening when the `diagonal` is of `rank=1` and `k[0] != k[1]` which needs to taken care. This is not problem with tensor type of `k` . 

You can try below code which works for `k `as tensor also.

```
import tensorflow as tf
tf.raw_ops.MatrixDiagV3(
    diagonal=tf.constant([10],dtype=tf.float32),
    k=tf.constant([1,1],shape=[2],dtype=tf.int32), // k=(1,1) also works
    num_rows=-1,
    num_cols=-1,
    padding_value=0.0,
    align='RIGHT_LEFT',
    name=None
)
```

Proposing a probable fix for this.
 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/62984">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/62984">No</a>

### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

tf 2.17

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 22.04 LTS

### Mobile device

_No response_

### Python version

3.11.7

### Bazel version

6.5.0

### GCC/compiler version

clang 16

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

C++ API `MatrixDiagV3` aborts with tensor type of `k`.
Usually `k` would be given as an array of integers. But it can also accept a tensor, and it can lead to crash.
Suspected code location is [here](https://github.com/tensorflow/tensorflow/blob/e193d8ea7776ef5c6f5d769b6fb9c070213e737a/tensorflow/core/kernels/linalg/matrix_diag_op.cc#L241),
But I have no idea why it is not triggered when it runs with an integer array `k`.

### Standalone code to reproduce the issue
Note: You may have to run it several times to see crash.

```C++
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"

using namespace tensorflow;

int main() {
  SessionOptions options;
  ConfigProto & config = options.config;
  config.set_inter_op_parallelism_threads(1);
  config.set_intra_op_parallelism_threads(1);
  config.set_use_per_session_threads(false);
  std::unique_ptr<tensorflow::Session>
    session(tensorflow::NewSession(options));
  Scope scope = Scope::NewRootScope();

  Input k = ops::RandomUniformInt(scope, {2}, 0, (1 << 15) - 1);
  auto target = ops::MatrixDiagV3(scope.WithOpName("target"), {0.0f}, k, 1, 1, 0.0f);

  GraphDef graph_def;
  TF_CHECK_OK(scope.ToGraphDef(&graph_def));

  Status status = session->Create(graph_def);
  if (!status.ok()) {
    LOG(WARNING) << "Could not create session: " << status.message();
  }

  std::vector<Tensor> outputs;
  status = session->Run({}, {"target"}, {"target"}, &outputs);
  if (!status.ok()) {
    LOG(WARNING) << "Could not run session: " << status.message();
  }

  return 0;
}
```


### Relevant log output

```shell
2024-02-18 02:54:37.171134: F tensorflow/core/framework/tensor_shape.cc:356] Check failed: d >= 0 (0 vs. -1)
Aborted (core dumped)
```
