def _register_checkpoint_pytrees():
    'Registers TF custom container types as pytrees.'
    m = tf.Module()
    m.a = (tf.Module(), tf.Module())
    m.b = [tf.Module(), tf.Module()]
    m.c = {'a': tf.Module()}
    tuple_wrapper = type(m.a)
    list_wrapper = type(m.b)
    dict_wrapper = type(m.c)
    assert (tuple_wrapper is not tuple)
    assert (list_wrapper is not list)
    assert (dict_wrapper is not dict)
    jax.tree_util.register_pytree_node(tuple_wrapper, (lambda xs: (tuple(xs), None)), (lambda _, xs: tuple(xs)))
    jax.tree_util.register_pytree_node(list_wrapper, (lambda xs: (tuple(xs), None)), (lambda _, xs: list(xs)))
    jax.tree_util.register_pytree_node(dict_wrapper, (lambda s: (tuple(s.values()), tuple(s.keys()))), (lambda k, xs: dict(zip(k, xs))))
Registers TF custom container types as pytrees.
def _register_checkpoint_pytrees():
    
    m = tf.Module()
    m.a = (tf.Module(), tf.Module())
    m.b = [tf.Module(), tf.Module()]
    m.c = {'a': tf.Module()}
    tuple_wrapper = type(m.a)
    list_wrapper = type(m.b)
    dict_wrapper = type(m.c)
    assert (tuple_wrapper is not tuple)
    assert (list_wrapper is not list)
    assert (dict_wrapper is not dict)
    jax.tree_util.register_pytree_node(tuple_wrapper, (lambda xs: (tuple(xs), None)), (lambda _, xs: tuple(xs)))
    jax.tree_util.register_pytree_node(list_wrapper, (lambda xs: (tuple(xs), None)), (lambda _, xs: list(xs)))
    jax.tree_util.register_pytree_node(dict_wrapper, (lambda s: (tuple(s.values()), tuple(s.keys()))), (lambda k, xs: dict(zip(k, xs))))