Crash in `ResourceSparseApplyProximalAdagrad`

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

2.20.0-dev20250516

Custom code

Yes

OS platform and distribution

Linux (Docker container)

Mobile device

No response

Python version

3.12

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

TensorFlow crashes with a fatal error Check failed: d < dims() (1 vs. 1) when using tf.raw_ops.ResourceSparseApplyProximalAdagrad with large float32 values.

The crash occurs in tensorflow/core/framework/tensor_shape.cc:359 and results in an aborted process with core dump.

Should raise a Python exception with a descriptive error message.

The issue can be reproduced in this Colab notebook: https://colab.research.google.com/drive/1HoC35YNBUC-Fs_6vOgrimGoLxqBisKHr?usp=sharing

Standalone code to reproduce the issue

import tensorflow as tf
import numpy as np

# Create resource variables with extreme values
large_val = 5.24393461e+36

var = tf.Variable([[large_val, large_val, large_val],
                   [large_val, large_val, large_val]], 
                  dtype=tf.float32, name="var")

accum = tf.Variable([[large_val, large_val, large_val],
                     [large_val, large_val, large_val]], 
                    dtype=tf.float32, name="accum")

# Scalar tensors with large values
lr = tf.constant(large_val, dtype=tf.float32)
l1 = tf.constant(large_val, dtype=tf.float32) 
l2 = tf.constant(large_val, dtype=tf.float32)

# Gradient and indices
grad = tf.constant([7.90505e+31], dtype=tf.float32)
indices = tf.constant([0], dtype=tf.int32)

# This call causes the crash
result = tf.raw_ops.ResourceSparseApplyProximalAdagrad(
    var=var.handle,
    accum=accum.handle,
    lr=lr,
    l1=l1,
    l2=l2,
    grad=grad,
    indices=indices,
    use_locking=False
)

Relevant log output

The program crashes with a fatal error:

2025-05-24 17:38:59.747282: F tensorflow/core/framework/tensor_shape.cc:359] Check failed: d < dims() (1 vs. 1)
Aborted (core dumped)


Complete Log:


2025-05-24 17:38:58.413411: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2025-05-24 17:38:58.458995: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2025-05-24 17:38:59.456168: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
TensorFlow version: 2.20.0-dev20250516
2025-05-24 17:38:59.672876: E external/local_xla/xla/stream_executor/cuda/cuda_platform.cc:51] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
Tensor shapes:
var: (2, 3)
accum: (2, 3)
grad: (1,)
indices: (1,)
Tensor values:
var values: [[5.2439346e+36 5.2439346e+36 5.2439346e+36]
 [5.2439346e+36 5.2439346e+36 5.2439346e+36]]
lr: 5.2439346057351246e+36  
grad: [7.90505e+31]
2025-05-24 17:38:59.747282: F tensorflow/core/framework/tensor_shape.cc:359] Check failed: d < dims() (1 vs. 1)
Aborted (core dumped)