keras/src/trainers/data_adapters/array_slicing.py/JaxSliceable/convert_to_numpy
class JaxSliceable:    def convert_to_numpy(cls, x):
        from keras.src.backend.jax.core import convert_to_numpy

        return convert_to_numpy(x)