Package: refx
RefxAreference systemfor jaxWhat is Refx?Refx defines a simple reference system that allows you to create DAGs on top of JAX PyTrees. It enables some key features:Shared stateTractable mutabilitySemantic partitioningRefx is intended to be a low-level library that can be used as a building block within other JAX libraries.Why Refx?Functional systems likeflaxandhaikuare powerful but add a lot of complexity that is often transfered to the user. On the other hand, pytree-based systems likeequinoxare simpler but lack the ability to share parameters and modules.Refx aims to create a system that can be used to build neural networks libraries that has the simplicity of pytree-based systems while also having the power of functional systems.InstallationpipinstallrefxGetting StartedRefx's main data structure is theRefclass. It is a wrapper around a value that can be used as leaves in a pytree. It also has avalueattribute that can be used to access and mutate the value.importjaximportrefxr1=refx.Ref(1)r2=refx.Ref(2)pytree={'a':[r1,r1,r2],'b':r2}pytree['a'][0].value=10assertpytree['a'][1].value==10Refis not a pytree node, therefore you cannot pass pytrees containingRefs to JAX functions. To interact with JAX,refxprovides the following functions:deref: converts a pytree of references to a pytree of values and indexes.reref: converts a pytree of values and indexes to a pytree of references.derefmust be called before crossing a JAX boundary andrerefmust be called after crossing a JAX boundary.pytree=refx.deref(pytree)@jax.jitdeff(pytree):pytree=refx.reref(pytree)# perform some computation / update the referencespytree['a'][2].value=50returnjax.deref(pytree)pytree=f(pytree)pytree=refx.reref(pytree)assertpytree['b'].value==50As you see in the is example, we've effectively implemented shared state and tracktable mutability with pure pytrees.Trace-level awarenessIn JAX, unconstrained mutability can lead to tracer leakage. To prevent this,refxonly allows mutating references from the same trace level they were created in.r=refx.Ref(1)@jax.jitdeff():# ValueError: Cannot mutate ref from different trace levelr.value=jnp.array(1.0)...PartitioningEach reference has acollection: Hashableattribute that can be used to partition references into different groups.refxprovides thetree_partitionand to partition a pytree based a predicate function.r1=refx.Ref(1,collection="params")r2=refx.Ref(2,collection="batch_stats")pytree={'a':[r1,r1,r2],'b':r2}(params,rest),treedef=refx.tree_partition(pytree,lambdax:isinstance(x,refx.Ref)andx.collection=="params")assertparams=={('a','0'):r1,('a','1'):r1,('a','2'):refx.NOTHING,('b',):refx.NOTHING}assertrest=={('a','0'):refx.NOTHING,('a','1'):refx.NOTHING,('a','2'):r2,('b',):r2,}You can use more partitioning functions to partition a pytree into multiple groups,tree_partitionwill always return one more partition than the number of functions passed to it. The last partition (rest) will contain all the remaining elements of the pytree.tree_partitionalso returns atreedefthat can be used to reconstruct the pytree by using themerge_partitionsfunction:pytree=refx.merge_partitions(partitions,treedef)If you only need a single partition, you can use theget_partitionfunction:r1=refx.Ref(1,collection="params")r2=refx.Ref(2,collection="batch_stats")pytree={'a':[r1,r1,r2],'b':r2}params=refx.get_partition(pytree,lambdax:isinstance(x,refx.Ref)andx.collection=="params")assertparams=={('a','0'):r1,('a','1'):r1,('a','2'):refx.NOTHING,('b',):refx.NOTHING}Updating referencesrefxprovides theupdate_refsfunction to update references in a pytree. It updates atargetpytree with references reference with the values from asourcepytree. The source pytree can be either a pytree of references or a pytree of values. As an example, here is how you can useupdate_refsto perform gradient descent on a pytree of references:defby_collection(c):returnlambdax:isinstance(x,refx.Ref)andx.collection==c(params,rest),treedef=refx.tree_partition(refx.deref(pytree),by_collection("params"))defloss(params):pytree=refx.merge_partitions((params,rest),treedef)...grads=jax.grad(loss)(params)# gradient descentparams=jax.tree_map(lambdaw,g:w-0.1*g,params,grads)refx.update_refs(refx.get_partition(pytree,by_collection("params")),params)