I'm trying to extend keras.ops.searchsorted
to work on 2D tensors. Natively it supports 1D only, and the error message suggests to use vectorized_map. However this throws an error because searchsorted tries to call len
on a symbolic tensor instead of shape.
Minimal example:
import keras
xs = keras.ops.stack([keras.ops.linspace(0, 1, 5) for _ in range(20)], axis=-1)
inputs = keras.random.uniform((100, 20), seed=42)
# this is okay
keras.ops.searchsorted(xs[:, 0], inputs[:, 0])
# this thwors
keras.ops.vectorized_map(
lambda i: keras.ops.searchsorted(xs[:, i], inputs[:, i]),
keras.ops.arange(20)
)
Throws:
TypeError: in user code:
File "/tmp/ipython-input-416-1753265487.py", line 2, in <lambda>
keras.ops.vectorized_map(lambda i: keras.ops.searchsorted(xs[:, i], inputs[:, i]), keras.ops.arange(20)
File ".../py/keras/src/ops/numpy.py", line 5278, in searchsorted
return backend.numpy.searchsorted(sorted_sequence, values, side=side)
File ".../keras/src/backend/tensorflow/numpy.py", line 2103, in searchsorted
"int32" if len(sorted_sequence) <= np.iinfo(np.int32).max else "int64"
TypeError: len is not well defined for a symbolic Tensor (loop_body/strided_slice:0). Please call `x.shape` rather than `len(x)` for shape information.
Comment From: iezepov
Note that this is true for TF backend. With jax backend it just works somehow.