`tf.raw_ops.ArgMax`: Heap buffer overflow
@Sehun0819 Could you please have a look at the [gist](https://colab.research.google.com/gist/sushreebarsa/ccbad423d1f928fdf3f24519ba5f9153/63070.ipynb) and confirm the issue?
Thank you! @sushreebarsa 
Even it does not print out ASAN log(as it wasn't built with ASAN flags), the error message seems to be an evidence.
`Expected dimension in the range [-4, 4), but got 1945239553`
`1945239553` is `0x73F20001` in hexadecimal. I guess the right part, `0x0001` came from 2byte input argument `tf.constant(1,shape=[],dtype=tf.int16)` and the left part `0x73F2` is the 2bytes that was not meant to be read. @Sehun0819 , Thanks for reporting. Proposed a fix which might fix this issue.  @SuryanarayanaY
Thank you!
There is one more code we should handle, [here](https://github.com/tensorflow/tensorflow/blob/50bfa12ba07bcd797be1d3841d059b4418a565de/tensorflow/core/ops/math_ops.cc#L1125-L1130):
```C++
  int64_t dimension_val;
  if (dim_t->dtype() == DT_INT32) {
    dimension_val = dim_t->scalar<int32>()();
  } else {
    dimension_val = dim_t->scalar<int64_t>()();
  }
```
This leads another heap buffer overflow, because it just assumes int32 and int64.

Reproduction:
```Python
import tensorflow as tf

tf.compat.v1.disable_eager_execution()

tf.raw_ops.ArgMax(
    input=tf.random.normal([1,1,1,1]),
    dimension=tf.constant(1,shape=[],dtype=tf.int16),
    output_type=tf.dtypes.int64,
    name=None
)
```
Log:
```bash
=================================================================
==3663542==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x6090000000c0 at pc 0x7f1245ed2a73 bp 0x7ffd5a214450 sp 0x7ffd5a214448
READ of size 8 at 0x6090000000c0 thread T0
    #0 0x7f1245ed2a72 in tensorflow::(anonymous namespace)::ArgOpShape(tensorflow::shape_inference::InferenceContext*) /proc/self/cwd/tensorflow/core/ops/math_ops.cc:1129:21
    #1 0x7f124bd6dd80 in absl::lts_20230802::Status std::__invoke_impl<absl::lts_20230802::Status, absl::lts_20230802::Status (*&)(tensorflow::shape_inference::InferenceContext*), tensorflow::shape_inference::InferenceContext*>(std::__invoke_other, absl::lts_20230802::Status (*&)(tensorflow::shape_inference::InferenceContext*), tensorflow::shape_inference::InferenceContext*&&) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/invoke.h:61:14
    #2 0x7f124bd6dd80 in std::enable_if<is_invocable_r_v<absl::lts_20230802::Status, absl::lts_20230802::Status (*&)(tensorflow::shape_inference::InferenceContext*), tensorflow::shape_inference::InferenceContext*>, absl::lts_20230802::Status>::type std::__invoke_r<absl::lts_20230802::Status, absl::lts_20230802::Status (*&)(tensorflow::shape_inference::InferenceContext*), tensorflow::shape_inference::InferenceContext*>(absl::lts_20230802::Status (*&)(tensorflow::shape_inference::InferenceContext*), tensorflow::shape_inference::InferenceContext*&&) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/invoke.h:114:9
    #3 0x7f124bd6dd80 in std::_Function_handler<absl::lts_20230802::Status (tensorflow::shape_inference::InferenceContext*), absl::lts_20230802::Status (*)(tensorflow::shape_inference::InferenceContext*)>::_M_invoke(std::_Any_data const&, tensorflow::shape_inference::InferenceContext*&&) /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/std_function.h:290:9
    #4 0x7f124c5b4083 in std::function<absl::lts_20230802::Status (tensorflow::shape_inference::InferenceContext*)>::operator()(tensorflow::shape_inference::InferenceContext*) const /usr/lib/gcc/x86_64-linux-gnu/11/../../../../include/c++/11/bits/std_function.h:590:9
    #5 0x7f124c5b4083 in tensorflow::shape_inference::InferenceContext::Run(std::function<absl::lts_20230802::Status (tensorflow::shape_inference::InferenceContext*)> const&) /proc/self/cwd/tensorflow/core/framework/shape_inference.cc:106:14
    #6 0x7f124bd65c0b in tensorflow::ShapeRefiner::RunShapeFn(tensorflow::Node const*, tensorflow::OpRegistrationData const*, tensorflow::ExtendedInferenceContext*, tensorflow::shape_inference::InferenceContext*)::$_0::operator()() const /proc/self/cwd/tensorflow/core/common_runtime/shape_refiner.cc:760:7
    #7 0x7f124bd5d360 in tensorflow::ShapeRefiner::RunShapeFn(tensorflow::Node const*, tensorflow::OpRegistrationData const*, tensorflow::ExtendedInferenceContext*, tensorflow::shape_inference::InferenceContext*) /proc/self/cwd/tensorflow/core/common_runtime/shape_refiner.cc:827:7
    #8 0x7f124bd5a54d in tensorflow::ShapeRefiner::AddNodeInternal(tensorflow::Node const*, tensorflow::shape_inference::InferenceContext*) /proc/self/cwd/tensorflow/core/common_runtime/shape_refiner.cc:270:3
    #9 0x7f124bd5c6df in tensorflow::ShapeRefiner::AddNode(tensorflow::Node const*) /proc/self/cwd/tensorflow/core/common_runtime/shape_refiner.cc:212:10
    #10 0x7f121547e005 in TF_FinishOperationLocked /proc/self/cwd/tensorflow/c/c_api.cc:1063:45
    #11 0x7f121547e5f0 in TF_FinishOperation /proc/self/cwd/tensorflow/c/c_api.cc:1082:10
    #12 0x7f11f70af947 in pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52::operator()(TF_OperationDescription*) const /proc/self/cwd/tensorflow/python/client/tf_session_wrapper.cc:1514:32
    #13 0x7f11f70af947 in TF_Operation* pybind11::detail::argument_loader<TF_OperationDescription*>::call_impl<TF_Operation*, pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&, 0ul, pybind11::detail::void_type>(pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&, std::integer_sequence<unsigned long, 0ul>, pybind11::detail::void_type&&) && /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/cast.h:1443:16
    #14 0x7f11f70af947 in std::enable_if<!std::is_void<TF_Operation*>::value, TF_Operation*>::type pybind11::detail::argument_loader<TF_OperationDescription*>::call<TF_Operation*, pybind11::detail::void_type, pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&>(pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&) && /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/cast.h:1411:42
    #15 0x7f11f70af947 in void pybind11::cpp_function::initialize<pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52, TF_Operation*, TF_OperationDescription*, pybind11::name, pybind11::scope, pybind11::sibling, pybind11::return_value_policy>(pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&&, TF_Operation* (*)(TF_OperationDescription*), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&, pybind11::return_value_policy const&)::'lambda'(pybind11::detail::function_call&)::operator()(pybind11::detail::function_call&) const /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/pybind11.h:248:69
    #16 0x7f11f70af947 in void pybind11::cpp_function::initialize<pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52, TF_Operation*, TF_OperationDescription*, pybind11::name, pybind11::scope, pybind11::sibling, pybind11::return_value_policy>(pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&&, TF_Operation* (*)(TF_OperationDescription*), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&, pybind11::return_value_policy const&)::'lambda'(pybind11::detail::function_call&)::__invoke(pybind11::detail::function_call&) /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/pybind11.h:223:21
    #17 0x7f11f70ea2a4 in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/pybind11.h:939:30
    #18 0x528186 in cfunction_call /usr/local/src/conda/python-3.11.7/Objects/methodobject.c:542:18
    #19 0x503a0b in _PyObject_MakeTpCall /usr/local/src/conda/python-3.11.7/Objects/call.c:214:18
    #20 0x510f32 in _PyEval_EvalFrameDefault /usr/local/src/conda/python-3.11.7/Python/ceval.c:4769:23
    #21 0x538732 in _PyEval_EvalFrame /usr/local/src/conda/python-3.11.7/Include/internal/pycore_ceval.h:73:16
    #22 0x538732 in _PyEval_Vector /usr/local/src/conda/python-3.11.7/Python/ceval.c:6434:24
    #23 0x538732 in _PyFunction_Vectorcall /usr/local/src/conda/python-3.11.7/Objects/call.c:393:16
    #24 0x5426bb in _PyVectorcall_Call /usr/local/src/conda/python-3.11.7/Objects/call.c:257:24
    #25 0x5426bb in _PyObject_Call /usr/local/src/conda/python-3.11.7/Objects/call.c:328:16
    #26 0x5426bb in PyObject_Call /usr/local/src/conda/python-3.11.7/Objects/call.c:355:12
    #27 0x514ff0 in do_call_core /usr/local/src/conda/python-3.11.7/Python/ceval.c:7352:12
    #28 0x514ff0 in _PyEval_EvalFrameDefault /usr/local/src/conda/python-3.11.7/Python/ceval.c:5376:22
    #29 0x538732 in _PyEval_EvalFrame /usr/local/src/conda/python-3.11.7/Include/internal/pycore_ceval.h:73:16
    #30 0x538732 in _PyEval_Vector /usr/local/src/conda/python-3.11.7/Python/ceval.c:6434:24
    #31 0x538732 in _PyFunction_Vectorcall /usr/local/src/conda/python-3.11.7/Objects/call.c:393:16
    #32 0x5426bb in _PyVectorcall_Call /usr/local/src/conda/python-3.11.7/Objects/call.c:257:24
    #33 0x5426bb in _PyObject_Call /usr/local/src/conda/python-3.11.7/Objects/call.c:328:16
    #34 0x5426bb in PyObject_Call /usr/local/src/conda/python-3.11.7/Objects/call.c:355:12
    #35 0x514ff0 in do_call_core /usr/local/src/conda/python-3.11.7/Python/ceval.c:7352:12
    #36 0x514ff0 in _PyEval_EvalFrameDefault /usr/local/src/conda/python-3.11.7/Python/ceval.c:5376:22
    #37 0x5cb559 in _PyEval_EvalFrame /usr/local/src/conda/python-3.11.7/Include/internal/pycore_ceval.h:73:16
    #38 0x5cb559 in _PyEval_Vector /usr/local/src/conda/python-3.11.7/Python/ceval.c:6434:24
    #39 0x5cac2e in PyEval_EvalCode /usr/local/src/conda/python-3.11.7/Python/ceval.c:1148:21
    #40 0x5ebcf6 in run_eval_code_obj /usr/local/src/conda/python-3.11.7/Python/pythonrun.c:1710:9
    #41 0x5e788f in run_mod /usr/local/src/conda/python-3.11.7/Python/pythonrun.c:1731:19
    #42 0x5fc831 in pyrun_file /usr/local/src/conda/python-3.11.7/Python/pythonrun.c:1626:15
    #43 0x5fbbfe in _PyRun_SimpleFileObject /usr/local/src/conda/python-3.11.7/Python/pythonrun.c:440:13
    #44 0x5fb922 in _PyRun_AnyFileObject /usr/local/src/conda/python-3.11.7/Python/pythonrun.c:79:15
    #45 0x5f65cd in pymain_run_file_obj /usr/local/src/conda/python-3.11.7/Modules/main.c:360:15
    #46 0x5f65cd in pymain_run_file /usr/local/src/conda/python-3.11.7/Modules/main.c:379:15
    #47 0x5f65cd in pymain_run_python /usr/local/src/conda/python-3.11.7/Modules/main.c:601:21
    #48 0x5f65cd in Py_RunMain /usr/local/src/conda/python-3.11.7/Modules/main.c:680:5
    #49 0x5bb3d8 in Py_BytesMain /usr/local/src/conda/python-3.11.7/Modules/main.c:734:12
    #50 0x7f1321691d8f in __libc_start_call_main csu/../sysdeps/nptl/libc_start_call_main.h:58:16
    #51 0x7f1321691e3f in __libc_start_main csu/../csu/libc-start.c:392:3
    #52 0x5bb222 in _start (/home/loft/anaconda3/envs/tf-latest-asan/bin/python3.11+0x5bb222)

0x6090000000c2 is located 0 bytes after 2-byte region [0x6090000000c0,0x6090000000c2)
allocated by thread T0 here:
    #0 0x7f1321a76617 in __interceptor_posix_memalign /home/runner/work/llvm-project/llvm-project/final/llvm-project/compiler-rt/lib/asan/asan_malloc_linux.cpp:145:3
    #1 0x7f124e776902 in tsl::port::AlignedMalloc(unsigned long, int) (/home/loft/anaconda3/envs/tf-latest-asan/lib/python3.11/site-packages/tensorflow/python/platform/../../libtensorflow_framework.so.2+0x4a18902)
    #2 0x7f124c553bcc in tsl::(anonymous namespace)::CPUAllocator::AllocateRaw(unsigned long, unsigned long) cpu_allocator_impl.cc
    #3 0x7f124c71e19f in short* tensorflow::TypedAllocator::Allocate<short>(tsl::Allocator*, unsigned long, tsl::AllocationAttributes const&) /proc/self/cwd/./tensorflow/core/framework/typed_allocator.h:47:24
    #4 0x7f124c71e19f in tensorflow::(anonymous namespace)::Buffer<short>::Buffer(tsl::Allocator*, long) /proc/self/cwd/tensorflow/core/framework/tensor.cc:568:21
    #5 0x7f124c71e19f in tensorflow::TensorBuffer* tensorflow::(anonymous namespace)::FromProtoField<short>(tsl::Allocator*, tensorflow::TensorProto const&, long) /proc/self/cwd/tensorflow/core/framework/tensor.cc:600:24
    #6 0x7f124c70ddd1 in tensorflow::Tensor::FromProto(tsl::Allocator*, tensorflow::TensorProto const&) /proc/self/cwd/tensorflow/core/framework/tensor.cc:1140:7
    #7 0x7f124bcf4e5b in tensorflow::EvaluateConstantTensor(tensorflow::Node const&, int, tensorflow::ShapeRefiner const&, absl::lts_20230802::FunctionRef<std::optional<tensorflow::Tensor> (tensorflow::Node const&, int)>, std::optional<tensorflow::EvaluateConstantTensorRunner>) /proc/self/cwd/tensorflow/core/common_runtime/eval_const_tensor.cc:375:9
    #8 0x7f124bd613ab in tensorflow::ShapeRefiner::EvaluateConstantTensorForEdge(tensorflow::Node const*, int, bool*, tensorflow::Tensor*, tensorflow::shape_inference::InferenceContext*) /proc/self/cwd/tensorflow/core/common_runtime/shape_refiner.cc:454:3
    #9 0x7f124bd5cf06 in tensorflow::ShapeRefiner::RunShapeFn(tensorflow::Node const*, tensorflow::OpRegistrationData const*, tensorflow::ExtendedInferenceContext*, tensorflow::shape_inference::InferenceContext*) /proc/self/cwd/tensorflow/core/common_runtime/shape_refiner.cc:799:9
    #10 0x7f124bd5a54d in tensorflow::ShapeRefiner::AddNodeInternal(tensorflow::Node const*, tensorflow::shape_inference::InferenceContext*) /proc/self/cwd/tensorflow/core/common_runtime/shape_refiner.cc:270:3
    #11 0x7f124bd5c6df in tensorflow::ShapeRefiner::AddNode(tensorflow::Node const*) /proc/self/cwd/tensorflow/core/common_runtime/shape_refiner.cc:212:10
    #12 0x7f121547e005 in TF_FinishOperationLocked /proc/self/cwd/tensorflow/c/c_api.cc:1063:45
    #13 0x7f121547e5f0 in TF_FinishOperation /proc/self/cwd/tensorflow/c/c_api.cc:1082:10
    #14 0x7f11f70af947 in pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52::operator()(TF_OperationDescription*) const /proc/self/cwd/tensorflow/python/client/tf_session_wrapper.cc:1514:32
    #15 0x7f11f70af947 in TF_Operation* pybind11::detail::argument_loader<TF_OperationDescription*>::call_impl<TF_Operation*, pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&, 0ul, pybind11::detail::void_type>(pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&, std::integer_sequence<unsigned long, 0ul>, pybind11::detail::void_type&&) && /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/cast.h:1443:16
    #16 0x7f11f70af947 in std::enable_if<!std::is_void<TF_Operation*>::value, TF_Operation*>::type pybind11::detail::argument_loader<TF_OperationDescription*>::call<TF_Operation*, pybind11::detail::void_type, pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&>(pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&) && /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/cast.h:1411:42
    #17 0x7f11f70af947 in void pybind11::cpp_function::initialize<pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52, TF_Operation*, TF_OperationDescription*, pybind11::name, pybind11::scope, pybind11::sibling, pybind11::return_value_policy>(pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&&, TF_Operation* (*)(TF_OperationDescription*), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&, pybind11::return_value_policy const&)::'lambda'(pybind11::detail::function_call&)::operator()(pybind11::detail::function_call&) const /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/pybind11.h:248:69
    #18 0x7f11f70af947 in void pybind11::cpp_function::initialize<pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52, TF_Operation*, TF_OperationDescription*, pybind11::name, pybind11::scope, pybind11::sibling, pybind11::return_value_policy>(pybind11_init__pywrap_tf_session(pybind11::module_&)::$_52&&, TF_Operation* (*)(TF_OperationDescription*), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&, pybind11::return_value_policy const&)::'lambda'(pybind11::detail::function_call&)::__invoke(pybind11::detail::function_call&) /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/pybind11.h:223:21
    #19 0x7f11f70ea2a4 in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/pybind11.h:939:30
    #20 0x528186 in cfunction_call /usr/local/src/conda/python-3.11.7/Objects/methodobject.c:542:18

SUMMARY: AddressSanitizer: heap-buffer-overflow /proc/self/cwd/tensorflow/core/ops/math_ops.cc:1129:21 in tensorflow::(anonymous namespace)::ArgOpShape(tensorflow::shape_inference::InferenceContext*)
Shadow bytes around the buggy address:
  0x608ffffffe00: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
  0x608ffffffe80: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
  0x608fffffff00: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
  0x608fffffff80: 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00
  0x609000000000: fa fa fa fa fa fa fa fa 00 00 fa fa fa fa fa fa
=>0x609000000080: fa fa fa fa fa fa fa fa[02]fa fa fa fa fa fa fa
  0x609000000100: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000180: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000200: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000280: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000300: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
Shadow byte legend (one shadow byte represents 8 application bytes):
  Addressable:           00
  Partially addressable: 01 02 03 04 05 06 07 
  Heap left redzone:       fa
  Freed heap region:       fd
  Stack left redzone:      f1
  Stack mid redzone:       f2
  Stack right redzone:     f3
  Stack after return:      f5
  Stack use after scope:   f8
  Global redzone:          f9
  Global init order:       f6
  Poisoned by user:        f7
  Container overflow:      fc
  Array cookie:            ac
  Intra object redzone:    bb
  ASan internal:           fe
  Left alloca redzone:     ca
  Right alloca redzone:    cb
==3663542==ABORTING
```
### 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?

`tf.raw_ops.ArgMax` can lead to heap buffer overflow.
[Error location](https://github.com/tensorflow/tensorflow/blob/774b0c3e97b5ef60bfc9c54961347dd0bc3660a8/tensorflow/core/kernels/argmax_op.cc#L59):
```C++
    const int32_t dim = internal::SubtleMustCopy(dimension.scalar<int32>()());
```
It copies scalar value of `dimension` without checking exact type. Therefore when `dimension` is an `int16` it reads over the bound.
Note that `int16` is an allowed type for `dimension` according to [opdef](https://github.com/tensorflow/tensorflow/blob/41b93a8b310086f69aab6b6369d2af9d5178881d/tensorflow/core/ops/math_ops.cc#L1153-L1160):
```C++
REGISTER_OP("ArgMax")
    .Input("input: T")
    .Input("dimension: Tidx")
    .Output("output: output_type")
    .Attr("T: {realnumbertype, quantizedtype, bool}")
    .Attr("Tidx: {int16, int32, int64} = DT_INT32")
    .Attr("output_type: {int16, uint16, int32, int64} = DT_INT64")
    .SetShapeFn(ArgOpShape);
```

### Standalone code to reproduce the issue

```python
import tensorflow as tf

tf.raw_ops.ArgMax(
    input=tf.random.normal([1,1,1,1]),
    dimension=tf.constant(1,shape=[],dtype=tf.int16),
    output_type=tf.dtypes.int64,
    name=None
)
```


### Relevant log output

The below log needs ASAN build.
```shell
=================================================================
==4008222==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x609000000400 at pc 0x7fa6a0dca809 bp 0x7ffe63b29d90 sp 0x7ffe63b29d88
READ of size 4 at 0x609000000400 thread T0
    #0 0x7fa6a0dca808 in int const tensorflow::internal::SubtleMustCopy<int>(int const&) /proc/self/cwd/./tensorflow/core/framework/bounds_check.h:49:10
    #1 0x7fa6a0dca808 in tensorflow::ArgOp<Eigen::ThreadPoolDevice, float, long, tensorflow::functor::ArgMax<Eigen::ThreadPoolDevice, float, long>>::Compute(tensorflow::OpKernelContext*) /proc/self/cwd/tensorflow/core/kernels/argmax_op.cc:59:25
    #2 0x7fa6bbbfab9f in tensorflow::ThreadPoolDevice::Compute(tensorflow::OpKernel*, tensorflow::OpKernelContext*) /proc/self/cwd/tensorflow/core/common_runtime/threadpool_device.cc:185:14
    #3 0x7fa6bb91644f in tensorflow::(anonymous namespace)::SingleThreadedExecutorImpl::Run(tensorflow::Executor::Args const&) /proc/self/cwd/tensorflow/core/common_runtime/single_threaded_executor.cc:445:15
    #4 0x7fa6bb836ae6 in tensorflow::FunctionLibraryRuntimeImpl::RunSync(tensorflow::FunctionLibraryRuntime::Options, unsigned long, absl::lts_20230802::Span<tensorflow::Tensor const>, std::vector<tensorflow::Tensor, std::allocator<tensorflow::Tensor>>*) /proc/self/cwd/tensorflow/core/common_runtime/function.cc:1339:3
    #5 0x7fa6bb866863 in tensorflow::ProcessFunctionLibraryRuntime::RunMultiDeviceSync(tensorflow::FunctionLibraryRuntime::Options const&, unsigned long, std::vector<std::variant<tensorflow::Tensor, tensorflow::TensorShape>, std::allocator<std::variant<tensorflow::Tensor, tensorflow::TensorShape>>>*, std::function<absl::lts_20230802::Status (tensorflow::ProcessFunctionLibraryRuntime::ComponentFunctionData const&, tensorflow::ProcessFunctionLibraryRuntime::InternalArgs*)>) const /proc/self/cwd/tensorflow/core/common_runtime/process_function_library_runtime.cc:931:16
    #6 0x7fa6bb877e9b in tensorflow::ProcessFunctionLibraryRuntime::RunSync(tensorflow::FunctionLibraryRuntime::Options const&, unsigned long, absl::lts_20230802::Span<tensorflow::Tensor const>, std::vector<tensorflow::Tensor, std::allocator<tensorflow::Tensor>>*) const /proc/self/cwd/tensorflow/core/common_runtime/process_function_library_runtime.cc:1526:21
    #7 0x7fa69db93761 in tensorflow::KernelAndDeviceFunc::Run(tensorflow::ScopedStepContainer*, tensorflow::EagerKernelArgs const&, std::vector<std::variant<tensorflow::Tensor, tensorflow::TensorShape>, std::allocator<std::variant<tensorflow::Tensor, tensorflow::TensorShape>>>*, tsl::CancellationManager*, std::optional<tensorflow::EagerFunctionParams> const&, std::optional<tensorflow::ManagedStackTrace> const&, tsl::CoordinationServiceAgent*) /proc/self/cwd/tensorflow/core/common_runtime/eager/kernel_and_device.cc:464:21
    #8 0x7fa69da529f5 in tensorflow::EagerKernelExecute(tensorflow::EagerContext*, absl::lts_20230802::InlinedVector<tensorflow::TensorHandle*, 4ul, std::allocator<tensorflow::TensorHandle*>> const&, std::optional<tensorflow::EagerFunctionParams> const&, tsl::core::RefCountPtr<tensorflow::KernelAndDevice> const&, tensorflow::GraphCollector*, tsl::CancellationManager*, absl::lts_20230802::Span<tensorflow::TensorHandle*>, std::optional<tensorflow::ManagedStackTrace> const&) /proc/self/cwd/tensorflow/core/common_runtime/eager/execute.cc:2167:3
    #9 0x7fa69da86630 in tensorflow::ExecuteNode::Run() /proc/self/cwd/./tensorflow/core/common_runtime/eager/execute_node.h:127:12
    #10 0x7fa69db7f849 in tensorflow::EagerExecutor::SyncExecute(tensorflow::EagerNode*) /proc/self/cwd/tensorflow/core/common_runtime/eager/eager_executor.cc:128:13
    #11 0x7fa69da512d9 in tensorflow::(anonymous namespace)::AddOrExecuteNode(tsl::core::RefCountPtr<tensorflow::KernelAndDevice>, tensorflow::EagerOperation*, tensorflow::TensorHandle**) /proc/self/cwd/tensorflow/core/common_runtime/eager/execute.cc:1685:25
    #12 0x7fa69da512d9 in tensorflow::(anonymous namespace)::EagerLocalExecute(tensorflow::EagerOperation*, tensorflow::TensorHandle**, int*) /proc/self/cwd/tensorflow/core/common_runtime/eager/execute.cc:1760:14
    #13 0x7fa69da4a107 in tensorflow::DoEagerExecute(tensorflow::EagerOperation*, tensorflow::TensorHandle**, int*) /proc/self/cwd/tensorflow/core/common_runtime/eager/execute.cc:2127:12
    #14 0x7fa69da54bd7 in tensorflow::EagerExecute(tensorflow::EagerOperation*, tensorflow::TensorHandle**, int*) /proc/self/cwd/tensorflow/core/common_runtime/eager/execute.cc:2207:10
    #15 0x7fa690dd1f99 in tensorflow::EagerOperation::Execute(absl::lts_20230802::Span<tensorflow::AbstractTensorHandle*>, int*) /proc/self/cwd/tensorflow/core/common_runtime/eager/core.cc:187:10
    #16 0x7fa69db7959e in tensorflow::CustomDeviceOpHandler::Execute(tensorflow::ImmediateExecutionOperation*, tensorflow::ImmediateExecutionTensorHandle**, int*) /proc/self/cwd/tensorflow/core/common_runtime/eager/custom_device_op_handler.cc
    #17 0x7fa6851e18f1 in TFE_Execute /proc/self/cwd/tensorflow/c/eager/c_api.cc:907:62
    #18 0x7fa6b989418d in TFE_Py_FastPathExecute_C(_object*) /proc/self/cwd/tensorflow/python/eager/pywrap_tfe_src.cc:3979:3
    #19 0x7fa667d2683e in pybind11_init__pywrap_tfe(pybind11::module_&)::$_60::operator()(pybind11::args) const /proc/self/cwd/tensorflow/python/tfe_wrapper.cc:1276:35
    #20 0x7fa667d2683e in pybind11::object pybind11::detail::argument_loader<pybind11::args>::call_impl<pybind11::object, pybind11_init__pywrap_tfe(pybind11::module_&)::$_60&, 0ul, pybind11::detail::void_type>(pybind11_init__pywrap_tfe(pybind11::module_&)::$_60&, std::integer_sequence<unsigned long, 0ul>, pybind11::detail::void_type&&) && /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/cast.h:1443:16
    #21 0x7fa667d2683e in std::enable_if<!std::is_void<pybind11::object>::value, pybind11::object>::type pybind11::detail::argument_loader<pybind11::args>::call<pybind11::object, pybind11::detail::void_type, pybind11_init__pywrap_tfe(pybind11::module_&)::$_60&>(pybind11_init__pywrap_tfe(pybind11::module_&)::$_60&) && /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/cast.h:1411:42
    #22 0x7fa667d2683e in void pybind11::cpp_function::initialize<pybind11_init__pywrap_tfe(pybind11::module_&)::$_60, pybind11::object, pybind11::args, pybind11::name, pybind11::scope, pybind11::sibling>(pybind11_init__pywrap_tfe(pybind11::module_&)::$_60&&, pybind11::object (*)(pybind11::args), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&)::'lambda'(pybind11::detail::function_call&)::operator()(pybind11::detail::function_call&) const /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/pybind11.h:248:69
    #23 0x7fa667d2683e in void pybind11::cpp_function::initialize<pybind11_init__pywrap_tfe(pybind11::module_&)::$_60, pybind11::object, pybind11::args, pybind11::name, pybind11::scope, pybind11::sibling>(pybind11_init__pywrap_tfe(pybind11::module_&)::$_60&&, pybind11::object (*)(pybind11::args), pybind11::name const&, pybind11::scope const&, pybind11::sibling const&)::'lambda'(pybind11::detail::function_call&)::__invoke(pybind11::detail::function_call&) /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/pybind11.h:223:21
    #24 0x7fa667d67a59 in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) /proc/self/cwd/bazel-out/k8-dbg/bin/external/pybind11/_virtual_includes/pybind11/pybind11/pybind11.h:939:30
    #25 0x528186 in cfunction_call /usr/local/src/conda/python-3.11.7/Objects/methodobject.c:542:18
    #26 0x503a0b in _PyObject_MakeTpCall /usr/local/src/conda/python-3.11.7/Objects/call.c:214:18
    #27 0x510f32 in _PyEval_EvalFrameDefault /usr/local/src/conda/python-3.11.7/Python/ceval.c:4769:23
    #28 0x538732 in _PyEval_EvalFrame /usr/local/src/conda/python-3.11.7/Include/internal/pycore_ceval.h:73:16
    #29 0x538732 in _PyEval_Vector /usr/local/src/conda/python-3.11.7/Python/ceval.c:6434:24
    #30 0x538732 in _PyFunction_Vectorcall /usr/local/src/conda/python-3.11.7/Objects/call.c:393:16
    #31 0x5426bb in _PyVectorcall_Call /usr/local/src/conda/python-3.11.7/Objects/call.c:257:24
    #32 0x5426bb in _PyObject_Call /usr/local/src/conda/python-3.11.7/Objects/call.c:328:16
    #33 0x5426bb in PyObject_Call /usr/local/src/conda/python-3.11.7/Objects/call.c:355:12
    #34 0x514ff0 in do_call_core /usr/local/src/conda/python-3.11.7/Python/ceval.c:7352:12
    #35 0x514ff0 in _PyEval_EvalFrameDefault /usr/local/src/conda/python-3.11.7/Python/ceval.c:5376:22
    #36 0x5cb559 in _PyEval_EvalFrame /usr/local/src/conda/python-3.11.7/Include/internal/pycore_ceval.h:73:16
    #37 0x5cb559 in _PyEval_Vector /usr/local/src/conda/python-3.11.7/Python/ceval.c:6434:24
    #38 0x5cac2e in PyEval_EvalCode /usr/local/src/conda/python-3.11.7/Python/ceval.c:1148:21
    #39 0x5ebcf6 in run_eval_code_obj /usr/local/src/conda/python-3.11.7/Python/pythonrun.c:1710:9
    #40 0x5e788f in run_mod /usr/local/src/conda/python-3.11.7/Python/pythonrun.c:1731:19
    #41 0x5fc831 in pyrun_file /usr/local/src/conda/python-3.11.7/Python/pythonrun.c:1626:15
    #42 0x5fbbfe in _PyRun_SimpleFileObject /usr/local/src/conda/python-3.11.7/Python/pythonrun.c:440:13
    #43 0x5fb922 in _PyRun_AnyFileObject /usr/local/src/conda/python-3.11.7/Python/pythonrun.c:79:15
    #44 0x5f65cd in pymain_run_file_obj /usr/local/src/conda/python-3.11.7/Modules/main.c:360:15
    #45 0x5f65cd in pymain_run_file /usr/local/src/conda/python-3.11.7/Modules/main.c:379:15
    #46 0x5f65cd in pymain_run_python /usr/local/src/conda/python-3.11.7/Modules/main.c:601:21
    #47 0x5f65cd in Py_RunMain /usr/local/src/conda/python-3.11.7/Modules/main.c:680:5
    #48 0x5bb3d8 in Py_BytesMain /usr/local/src/conda/python-3.11.7/Modules/main.c:734:12
    #49 0x7fa791408d8f in __libc_start_call_main csu/../sysdeps/nptl/libc_start_call_main.h:58:16
    #50 0x7fa791408e3f in __libc_start_main csu/../csu/libc-start.c:392:3
    #51 0x5bb222 in _start (/home/loft/anaconda3/envs/tf-latest-asan/bin/python3.11+0x5bb222)

0x609000000402 is located 0 bytes after 2-byte region [0x609000000400,0x609000000402)
allocated by thread T0 here:
    #0 0x7fa7917ed617 in __interceptor_posix_memalign /home/runner/work/llvm-project/llvm-project/final/llvm-project/compiler-rt/lib/asan/asan_malloc_linux.cpp:145:3
    #1 0x7fa6be4f6902 in tsl::port::AlignedMalloc(unsigned long, int) (/home/loft/anaconda3/envs/tf-latest-asan/lib/python3.11/site-packages/tensorflow/python/platform/../../libtensorflow_framework.so.2+0x4a18902)
    #2 0x7fa6bc2d3bcc in tsl::(anonymous namespace)::CPUAllocator::AllocateRaw(unsigned long, unsigned long) cpu_allocator_impl.cc
    #3 0x7fa6bc4883bc in short* tensorflow::TypedAllocator::Allocate<short>(tsl::Allocator*, unsigned long, tsl::AllocationAttributes const&) /proc/self/cwd/./tensorflow/core/framework/typed_allocator.h:47:24
    #4 0x7fa6bc4883bc in tensorflow::(anonymous namespace)::Buffer<short>::Buffer(tsl::Allocator*, long, tsl::AllocationAttributes const&) /proc/self/cwd/tensorflow/core/framework/tensor.cc:574:21
    #5 0x7fa6bc4883bc in tensorflow::Tensor::Tensor(tsl::Allocator*, tensorflow::DataType, tensorflow::TensorShape const&, tsl::AllocationAttributes const&) /proc/self/cwd/tensorflow/core/framework/tensor.cc:986:5
    #6 0x7fa6bbef871d in tensorflow::OpKernelContext::allocate_tensor(tensorflow::DataType, tensorflow::TensorShape const&, tensorflow::Tensor*, tsl::AllocatorAttributes, tsl::AllocationAttributes const&) /proc/self/cwd/tensorflow/core/framework/op_kernel.cc:764:10
    #7 0x7fa6bbef784e in tensorflow::OpKernelContext::allocate_tensor(tensorflow::DataType, tensorflow::TensorShape const&, tensorflow::Tensor*, tsl::AllocatorAttributes) /proc/self/cwd/./tensorflow/core/framework/op_kernel.h:1270:12
    #8 0x7fa6bbef784e in tensorflow::OpKernelContext::allocate_output(int, tensorflow::TensorShape const&, tensorflow::Tensor**, tsl::AllocatorAttributes) /proc/self/cwd/tensorflow/core/framework/op_kernel.cc:822:14
    #9 0x7fa6bbef4368 in tensorflow::OpKernelContext::allocate_output(int, tensorflow::TensorShape const&, tensorflow::Tensor**) /proc/self/cwd/tensorflow/core/framework/op_kernel.cc:728:10
    #10 0x7fa6a9bf0207 in tensorflow::CastOpBase::Compute(tensorflow::OpKernelContext*) (/home/loft/anaconda3/envs/tf-latest-asan/lib/python3.11/site-packages/tensorflow/python/platform/../../libtensorflow_cc.so.2+0x3e9fe207)
    #11 0x7fa6bbbfab9f in tensorflow::ThreadPoolDevice::Compute(tensorflow::OpKernel*, tensorflow::OpKernelContext*) /proc/self/cwd/tensorflow/core/common_runtime/threadpool_device.cc:185:14
    #12 0x7fa6bb91644f in tensorflow::(anonymous namespace)::SingleThreadedExecutorImpl::Run(tensorflow::Executor::Args const&) /proc/self/cwd/tensorflow/core/common_runtime/single_threaded_executor.cc:445:15
    #13 0x7fa6bb836ae6 in tensorflow::FunctionLibraryRuntimeImpl::RunSync(tensorflow::FunctionLibraryRuntime::Options, unsigned long, absl::lts_20230802::Span<tensorflow::Tensor const>, std::vector<tensorflow::Tensor, std::allocator<tensorflow::Tensor>>*) /proc/self/cwd/tensorflow/core/common_runtime/function.cc:1339:3
    #14 0x7fa6bb866863 in tensorflow::ProcessFunctionLibraryRuntime::RunMultiDeviceSync(tensorflow::FunctionLibraryRuntime::Options const&, unsigned long, std::vector<std::variant<tensorflow::Tensor, tensorflow::TensorShape>, std::allocator<std::variant<tensorflow::Tensor, tensorflow::TensorShape>>>*, std::function<absl::lts_20230802::Status (tensorflow::ProcessFunctionLibraryRuntime::ComponentFunctionData const&, tensorflow::ProcessFunctionLibraryRuntime::InternalArgs*)>) const /proc/self/cwd/tensorflow/core/common_runtime/process_function_library_runtime.cc:931:16
    #15 0x7fa6bb877e9b in tensorflow::ProcessFunctionLibraryRuntime::RunSync(tensorflow::FunctionLibraryRuntime::Options const&, unsigned long, absl::lts_20230802::Span<tensorflow::Tensor const>, std::vector<tensorflow::Tensor, std::allocator<tensorflow::Tensor>>*) const /proc/self/cwd/tensorflow/core/common_runtime/process_function_library_runtime.cc:1526:21
    #16 0x7fa69db93761 in tensorflow::KernelAndDeviceFunc::Run(tensorflow::ScopedStepContainer*, tensorflow::EagerKernelArgs const&, std::vector<std::variant<tensorflow::Tensor, tensorflow::TensorShape>, std::allocator<std::variant<tensorflow::Tensor, tensorflow::TensorShape>>>*, tsl::CancellationManager*, std::optional<tensorflow::EagerFunctionParams> const&, std::optional<tensorflow::ManagedStackTrace> const&, tsl::CoordinationServiceAgent*) /proc/self/cwd/tensorflow/core/common_runtime/eager/kernel_and_device.cc:464:21
    #17 0x7fa69da529f5 in tensorflow::EagerKernelExecute(tensorflow::EagerContext*, absl::lts_20230802::InlinedVector<tensorflow::TensorHandle*, 4ul, std::allocator<tensorflow::TensorHandle*>> const&, std::optional<tensorflow::EagerFunctionParams> const&, tsl::core::RefCountPtr<tensorflow::KernelAndDevice> const&, tensorflow::GraphCollector*, tsl::CancellationManager*, absl::lts_20230802::Span<tensorflow::TensorHandle*>, std::optional<tensorflow::ManagedStackTrace> const&) /proc/self/cwd/tensorflow/core/common_runtime/eager/execute.cc:2167:3
    #18 0x7fa69da86630 in tensorflow::ExecuteNode::Run() /proc/self/cwd/./tensorflow/core/common_runtime/eager/execute_node.h:127:12
    #19 0x7fa69db7f849 in tensorflow::EagerExecutor::SyncExecute(tensorflow::EagerNode*) /proc/self/cwd/tensorflow/core/common_runtime/eager/eager_executor.cc:128:13
    #20 0x7fa69da512d9 in tensorflow::(anonymous namespace)::AddOrExecuteNode(tsl::core::RefCountPtr<tensorflow::KernelAndDevice>, tensorflow::EagerOperation*, tensorflow::TensorHandle**) /proc/self/cwd/tensorflow/core/common_runtime/eager/execute.cc:1685:25
    #21 0x7fa69da512d9 in tensorflow::(anonymous namespace)::EagerLocalExecute(tensorflow::EagerOperation*, tensorflow::TensorHandle**, int*) /proc/self/cwd/tensorflow/core/common_runtime/eager/execute.cc:1760:14
    #22 0x7fa69da4a107 in tensorflow::DoEagerExecute(tensorflow::EagerOperation*, tensorflow::TensorHandle**, int*) /proc/self/cwd/tensorflow/core/common_runtime/eager/execute.cc:2127:12
    #23 0x7fa69da54bd7 in tensorflow::EagerExecute(tensorflow::EagerOperation*, tensorflow::TensorHandle**, int*) /proc/self/cwd/tensorflow/core/common_runtime/eager/execute.cc:2207:10
    #24 0x7fa690dd1f99 in tensorflow::EagerOperation::Execute(absl::lts_20230802::Span<tensorflow::AbstractTensorHandle*>, int*) /proc/self/cwd/tensorflow/core/common_runtime/eager/core.cc:187:10
    #25 0x7fa69db7959e in tensorflow::CustomDeviceOpHandler::Execute(tensorflow::ImmediateExecutionOperation*, tensorflow::ImmediateExecutionTensorHandle**, int*) /proc/self/cwd/tensorflow/core/common_runtime/eager/custom_device_op_handler.cc
    #26 0x7fa6851e18f1 in TFE_Execute /proc/self/cwd/tensorflow/c/eager/c_api.cc:907:62
    #27 0x7fa6b985dcbf in tensorflow::EagerCast(TFE_Context*, TFE_TensorHandle*, TF_DataType, TF_DataType, TSL_Status*) /proc/self/cwd/tensorflow/python/eager/pywrap_tensor.cc:259:3
    #28 0x7fa6b985e67b in tensorflow::ConvertToEagerTensorUncached(TFE_Context*, _object*, tensorflow::DataType, char const*) /proc/self/cwd/tensorflow/python/eager/pywrap_tensor.cc:317:11
    #29 0x7fa6b985f8a1 in tensorflow::ConvertToEagerTensor(TFE_Context*, _object*, tensorflow::DataType, char const*) /proc/self/cwd/tensorflow/python/eager/pywrap_tensor.cc:405:14
    #30 0x7fa6b9860088 in EagerTensor_init /proc/self/cwd/tensorflow/python/eager/pywrap_tensor.cc:529:18
    #31 0x5039d2 in type_call /usr/local/src/conda/python-3.11.7/Objects/typeobject.c:1103:19
    #32 0x5039d2 in _PyObject_MakeTpCall /usr/local/src/conda/python-3.11.7/Objects/call.c:214:18

SUMMARY: AddressSanitizer: heap-buffer-overflow /proc/self/cwd/./tensorflow/core/framework/bounds_check.h:49:10 in int const tensorflow::internal::SubtleMustCopy<int>(int const&)
Shadow bytes around the buggy address:
  0x609000000180: fa fa fa fa fa fa fa fa fd fa fa fa fa fa fa fa
  0x609000000200: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000280: fd fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000300: 04 fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000380: fd fd fd fd fa fa fa fa fa fa fa fa fa fa fa fa
=>0x609000000400:[02]fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000480: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000500: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000580: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000600: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
  0x609000000680: fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa fa
Shadow byte legend (one shadow byte represents 8 application bytes):
  Addressable:           00
  Partially addressable: 01 02 03 04 05 06 07 
  Heap left redzone:       fa
  Freed heap region:       fd
  Stack left redzone:      f1
  Stack mid redzone:       f2
  Stack right redzone:     f3
  Stack after return:      f5
  Stack use after scope:   f8
  Global redzone:          f9
  Global init order:       f6
  Poisoned by user:        f7
  Container overflow:      fc
  Array cookie:            ac
  Intra object redzone:    bb
  ASan internal:           fe
  Left alloca redzone:     ca
  Right alloca redzone:    cb
==4008222==ABORTING
```
