Update utils.py

serket/_src/utils.py

@@ -640,14 +640,4 @@ def map_func(view):
         result = jax.vmap(map_func)(views)
         return result.reshape(*output_shape, *result.shape[1:])
 
-    return single_call_wrapper
-
-
-def frozen_field(**kwargs):
-    """Freeze a field after setting it and unfreeze it before getting it."""
-    # this is useful for setting a field that is not a jax-type
-    # to allow the class to be passed across jax-boundaries]
-    return sk.field(
-        on_getattr=[*kwargs.pop("on_getattr", []), sk.unfreeze],
-        on_setattr=[*kwargs.pop("on_setattr", []), sk.freeze],
-    )
+    return single_call_wrapper