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

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.

Reproduce the call to Dataset.map with a jaxtyped function using beartype and fn_kwargs, then inspect datasets.arrow_dataset.Dataset._map_single and the emitted hashing warning. Compare behavior with and without the decorator; done means the iteration starts normally and dataset fingerprinting and caching no longer trigger the reported interruption or excessive delay.