Package: nux
NuX - Normalizing Flows using JAXWhat is NuX?NuX is a library for buildingnormalizing flowsusingJAX.What are normalizing flows?Normalizing flows learn a parametric model over an unknown probability density function using data.  We assume that data points are sampled i.i.d. from an unknown distribution p(x).  Normalizing flows learn a parametric approximation of the true data distribution, q(x), using maximum likelihood learning.  The resulting q(x) can be efficiently sampled from and evaluated exactly.Why use NuX?NuX has many normalizing flow layers implemented with an easy to use interface.importnux.flowsasnuximportjaxfromjaximportrandomimportjax.numpyasjnpkey=random.PRNGKey(0)# Build a dummy datasetx_train,x_test=jnp.ones((2,100,4))# Build a simple normalizing flowinit_fun=nux.sequential(nux.Coupling(),nux.ActNorm(),nux.UnitGaussianPrior())# Perform data-dependent initialization_,flow=init_fun(key,{'x':x_train},batched=True)# Run data through the flowinputs={'x':x_test}outputs,_=flow.apply(flow.params,flow.state,inputs)z,log_likelihood=outputs['x'],outputs['log_pz']+outputs['log_det']# Check the reconstructionsreconst,_=flow.apply(flow.params,flow.state,{'x':z},reverse=True)assertjnp.allclose(x_test,reconst['x'])What's implemented?Check out thebijective,injectiveandsurjectivetransformations that are implemented.  Any contributions are welcome!How does it work?The modularity of normalizing flows allows us to construct complex flows using code that is made simple using JAX.Create complex flowsFlow layers can be imported from nux.flows and chained together sequentially usingnux.flows.sequential.  For more complex flows, we can split the flow using theprobability chain ruleinnux.flows.ChainRuleand then run multiple flows in parallel usingnux.flows.parallel.Using these basic transformations, we can easily create complex flows.  For example, amultiscaleGLOWnormalizing flow can be implemented easily:importnux.flowsasnuxdefmulti_scale(flow,existing_flow):# This is implemented in nux.flows.composereturnnux.sequential(flow,nux.Squeeze(),nux.ChainRule(2,factor=True),nux.factored(existing_flow,nux.Identity()),nux.ChainRule(2,factor=False),nux.UnSqueeze())defGLOWBlock():returnnux.sequential(nux.ActNorm(),nux.OnebyOneConv(),nux.Coupling(n_channels=512))defGLOW(num_blocks=4):layers=[GLOWBlock()for_inrange(num_blocks)]returnnux.sequential(*layers)defMultiscaleGLOW(quantize_bits=3):flow=nux.Identity()flow=multi_scale(GLOW(),flow)flow=multi_scale(GLOW(),flow)flow=multi_scale(GLOW(),flow)flow=multi_scale(GLOW(),flow)flow=nux.sequential(nux.UniformDequantization(scale=2**quantize_bits),nux.Logit(),nux.Squeeze(),# So that the channel is divisible by 2flow,nux.Flatten(),nux.AffineGaussianPriorDiagCov(out_dim=128))# Use a low dimensional prior for best results!flow_init_fun=flow# The result of creating layers is an initializerreturnflow_init_funInitialize your flow with dataNuX initializes flows using data to infer the input and output shapes at each flow layer and to help initialize layers likeActnorm.  NuX uses dictionaries as the primary data-structure to pass data between flow layers.importjaximportjax.numpyasjnpflow_init_fun=MultiscaleGLOW()key=jax.random.PRNGKey(0)inputs={'x':jnp.zeros(64,32,32,3)}# Create a dummy datasetoutputs,flow=flow_init_fun(key,inputs,batched=True)# Must specify if the input data is batched or notInputs/OutputsEach flow application expects an input dictionary.  The key 'x' should correspond to the data that is passed between flow layers.  Furthermore all elements of the input dictionary are passed to each flow layer.  For example, in classification we can pass labels to any layer using:inputs={'x':data,'y':labels}Every flow layer returns a dictionary of outputs that contains the transformed data in key 'x' and a log likelihood contribution term.  For standard transformations, the log likelihood contribution term is under 'log_det' and 'log_pz' for priors.  Like the inputs, flow layers can also output other key value pairs.Flow data structureThe second value returned by an initializer call is theflow data structure.  This data structure contains the name of a layer and dictionaries of the input/output shapes/dims, parameters, state and apply function.  The shapes/dims aid auto-batching while the parameters and state parametrize the flow.  The difference between parameters and state is that parameters is intended to contain the parameters that will be trained with gradient descent while the state values do not (like the running statistics in batch normalization).jax.tree_util.tree_mapandjax.tree_util.tree_multimapare your friend when working with dictionaries!The apply function is called with the parameters, state, inputs, keyword arguments and a flag that specifies which direction to run the flow:# Run the flow forwards (x -> z)outputs,updated_state=flow.apply(flow.params,flow.state,inputs,key=key,reverse=False)log_px=outputs['log_pz']+outputs['log_det']# Run the flow forwards in reverse (z -> x)reconstr_inputs,_=flow.apply(flow.params,flow.state,outputs,key=key,reverse=True)UseHaikuto create deep flow layersFlow layers, likenux.flows.Coupling, can use a neural network to introduce complex non-linearities.  These neural networks must constructed usingHaiku.  There are default networks that flow layers default to, but any Haiku network can be used.  For example, we can construct a transformation for image coupling layers as follows:classSimpleConv(hk.Module):def__init__(self,out_shape,n_hidden_channels,name=None):super().__init__(name=name)_,_,out_channels=out_shapeself.out_channels=out_channelsself.n_hidden_channels=n_hidden_channelsself.last_channels=2*out_channelsdef__call__(self,x,**kwargs):H,W,C=x.shape# NuX ensures that the input will be unbatched!x=hk.Conv2D(output_channels=self.n_hidden_channels,kernel_shape=(3,3),stride=(1,1))(x[None])[0]x=jax.nn.relu(x)x=hk.Conv2D(output_channels=self.n_hidden_channels,kernel_shape=(1,1),stride=(1,1))(x[None])[0]x=jax.nn.relu(x)x=hk.Conv2D(output_channels=self.last_channels,kernel_shape=(3,3),stride=(1,1),w_init=hk.initializers.Constant(0),b_init=hk.initializers.Constant(0))(x[None])[0]mu,alpha=jnp.split(x,2,axis=-1)alpha=jnp.tanh(alpha)returnmu,alphaNuX handles batching internally, so every flow and network is guaranteed to be passed an unbatched input.Using JAXNuX is built using JAX so all of its features can be used on a flow.Creating custom flow layersCustom flows are easy to create in NuX. NuX can internally handle batch dimensions, so custom layers can be implemented assuming the input is unbatched:importjax.numpyasjnpimportjax.nn.initializersasjaxinitimportnux.flows.baseasbase@base.auto_batch# Ensure that apply_fun receives unbatched inputs.defOnebyOneConvDense(W_init=jaxinit.glorot_normal(),name='1x1conv_dense'):defapply_fun(params,state,inputs,reverse=False,**kwargs):x=inputs['x']# Unpack the inputsW=params['W']# Unpack the parametersheight,width,channel=x.shape# auto_batch ensures x is unbatched!# Compute the transformationif(reverse==False):z=jnp.einsum('ij,hwj->hwi',W,x)else:W_inv=jnp.linalg.inv(W)z=jnp.einsum('ij,hwj->hwi',W_inv,x)# Compute the log Jacobian determinantlog_det=jnp.linalg.slogdet(W)[1]log_det*=height*width# Return the outputs and update the state if necessary.outputs={'x':z,'log_det':log_det}updated_state=statereturnoutputs,updated_statedefcreate_params_and_state(key,input_shapes):# The data_dependent=False flag below ensures that input_shapes is unbatched.height,width,channel=input_shapes['x']# Initialize the parametersW=W_init(key,(channel,channel))params,state={'W':W},{}returnparams,statereturnbase.initialize(name,apply_fun,create_params_and_state,data_dependent=False)# Helper to put everything together.Under the hood,base.initializeextracts the shapes from the input to the initializer.  These shapes are passed tocreate_params_and_stateto generate the parameters and state.  The inputs, parameters and state are then passed toapply_funto compute the outputs.  Finally, the shapes/dimensions of the outputs are retrieved and stored.At runtime,base.auto_batchhas access to the unbatched input dimensions for each flow.  With this information, it recursively appliesjax.vmapto correctly handle nested batching.Thedata_dependentflag can be set to in order to pass the batched inputs tocreate_params_and_state.  ForActnorm, this looks like:@base.auto_batchdefActNorm(log_s_init=jaxinit.zeros,b_init=jaxinit.zeros,name='act_norm'):multiply_by=None# We can store initialize time constants in this outer scope.defapply_fun(params,state,inputs,reverse=False,**kwargs):if(reverse==False):z=(inputs['x']-params['b'])*jnp.exp(-params['log_s'])else:z=jnp.exp(params['log_s'])*inputs['x']+params['b']log_det=-params['log_s'].sum()*multiply_byoutputs={'x':z,'log_det':log_det}returnoutputs,statedefcreate_params_and_state(key,inputs,batch_depth):# The shape of x is the unbatched shape of x prepended with batch_depth dimensions.x=inputs['x']# Need to keep track of the dimensionality of all but the last axis in case we pass in an image.nonlocalmultiply_bymultiply_by=jnp.prod([sfori,sinenumerate(x.shape)ifi>=batch_depthandi<len(x.shape)-1])# Create the parameters using the batch of dataaxes=tuple(jnp.arange(len(x.shape)-1))params={'b':jnp.mean(x,axis=axes),'log_s':jnp.log(jnp.std(x,axis=axes)+1e-5)}state={}returnparams,statereturnbase.initialize(name,apply_fun,create_params_and_state,data_dependent=True)Testing a custom flownux.tests.nf_test.flow_testis a simple function to test the correctness of a flow.  It checks unbatched/batched/doubly-batched reconstructions by running a flow forwards then in reverse and checks the log Jacobian determinant against the brute force solution computed usingjax.jacobian.init_fun=Flow()unbatched_inputs={'x':data}flow_test(init_fun,unbatched_inputs,key)InstallationFor the moment, NuX only works with python 3.7.  The steps to install are:pip install nux
 pip install git+https://github.com/deepmind/dm-haikuIf you want GPU support for JAX, follow the intructions herehttps://github.com/google/jax#installation