nqxpack.registry.register_closure_simple_serialization

nqxpack.registry.register_closure_simple_serialization(parent_fun, closure_name, vars_mapping={}, original_qualname=None, override=False)

Register the closure resulting from a function to have a simple serialization.

This is usually enough for jax.nn.initializers, for example. The requirement for this to work is that the closure shows has a name of the form <parent_fun>.<locals>.<closure_name>.

To register, you must provide the parent function from which the closure was generated, the name of the closure, and optionally a mapping of the closure variable names to the names you want to use in the serialization.

Parameters:

• parent_fun – The parent function of the closure

• closure_name – The name of the closure

• vars_mapping – A mapping of the closure variable names to the names you want to use in the serialization

• original_qualname – The original qualname of the parent function. Jax usually hides the ._src part of the qualname, and you must provide it here to resolve the original function.

• override (bool)