keras/src/backend/tensorflow/numpy.py/wrapped
def _vmap_fn(fn, in_axes=0):
    if in_axes != 0:
        raise ValueError(
            "Not supported with `vectorize()` with the TensorFlow backend."
        )

    @functools.wraps(fn)
    def wrapped(x):
        return tf.vectorized_map(fn, x)

    return wrapped