liesel.liesel.goose.make_update_state_fn#

liesel.liesel.goose.make_update_state_fn(model, jaxify=True)[source]#

Returns a pure and jittable update_state() function for the provided model.

Return type

Callable[[NewType(Position, dict[str, Any]), dict[str, NodeState]], dict[str, NodeState]]