dataset.map(jaxtyped_function, fn_kwargs={"model": model}) will be very slow.

[jasonrichdarmawan]

dataset is Dataset from datasets

if we use @jaxtyped(typechecker=typechecker where typechecker is beartype, we need to press CTRL+C and get this error. Without this, the iteration will not start.

warnings.warn( ^CParameter 'function'=<function generate_entity_description_with_patching at 0x7740e09aa4d0> of the transform datasets.arrow_dataset.Dataset._map_single couldn't be hashed properly, a random hash was used instead. Make sure your transforms and parameters are serializable with pickle or dill for the dataset fingerprinting and caching to work. If you reuse this transform, the caching mechanism will consider it to be different from the previous calls and recompute everything. This warning is only showed once. Subsequent hashing failures won't be showed.

[patrick-kidger]

Can you provide a MWE without datasets?

FWIW I suspect this is a serialization issue – it sounds like they're trying to do some kind of multiprocessing and serializing wrapped function is always a tricky business.