Keras keras.ops.searchsorted doesn't work with symbolic tensors

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.