Package: jax-autovmap
Automatically broadcast via JAX vmapAutomatically broadcast a function that takes inputs of specific ranks to accept any batch dimensions usingjax.vmap.
See the documentation ofJAXfor more information aboutvmap.This module defines a decorator function which takes as input the base ranks of the arguments of a function.
The transformed function takes any broadcastable combination of batched inputs and automatically appliesjax.vmapas required.
One could equivalently applyvmapby hand, however the in-axes have to be chosen differently for different inputs.
The decorator takes care of this automatically; if the underlying ranks are known, batch dimensions can be inferred.ExampleConsider the following function which takes numeric arguments with fixed and known ranks as input:importjax.numpyasjnpdeffoo(s,v,m):# s - scalar, v - vector, m - matrixreturnv@m@v+s*v.sizeIf we have inputs of appropriate rank, the function can be applied without problem:s=jnp.array(2.0)# scalarv=jnp.ones(3)# vectorm=jnp.ones((3,3))# matrixprint(foo(s,v,m))# prints 15.0Assume now, however, that we have 5 matrices and 5 vectors for which we want to apply the above function as a batch.
Many numpy functions can take inputs with leading batch dimensions, but here we have an issue becausev @ m @ vrequiresmto be a matching matrix.s=jnp.array(2.0)v=jnp.ones((5,3))m=jnp.ones((5,3,3))print(foo(s,v,m))# throws TypeErrorThere are multiple possible ways we can solve thisWe could try to write the function more carefully, so it can take both single and batch inputs.
Or we could always require the first dimension to be a batch index (like often done in convolutional neural networks).Given known inputs, we can transform our function:jax.jit(foo, (None, 0, 0)).
However, the axes change depending on the inputs.
If we want to expose functions that accept batched inputs to the user, we need to have some clear naming scheme (to indicate which arguments are batched).Sometimes, the best solution is one of the above.
This module provides another more flexible solution.
If the ranks the functionwantsare known, we can derive which arguments have (leading) batch dimensions.
Based on that, we can applyjax.vmapappropriately.
That is exactly what theauto_vmapwrapper does.
Thanks tojax.jit, after the transformed function is JIT-compiled, there is no price to pay for this extra flexibility since it only depends on the statically-known input shapes.We can define the more flexible function as follows:fromjax_autovmapimportauto_vmap@auto_vmap(s=0,v=1,m=2)deffoo(s,v,m):returnv@m@v+s*v.sizeprint(foo(s,v,m))# prints [15. 15. 15. 15. 15.]The ranks can be specified by keyword argument as above, or positionally (in this case@auto_vmap(0, 1, 2)).
This does not have to be applied to all input arguments.
They can either be omitted or, if ranks are given positionally, specified asNone.If the arguments are pytrees (python structures of arrays) and the rank is a single integer, all constituents (leaves) are assumed to have that rank.
Alternatively, the rank can be a matching pytree, just like thein_axesinjax.vmap:@auto_vmap({'s'=0,'v'=1,'m'=2})deffoo(inputs):returninputs['v']@inputs['m']@inputs['v']print(foo(dict(s=s,v=v,m=m)))