Package: jax-jynx
JynxA straight forward neural network library written in jax. No hidden mechanisms, no black
magic. Requires only jax and optax.This library provides 3 components: (1) standard neural network layers, (2) afitfunction to train models, and (3) a collection of basic callbacks for
checkpointing and logging. Thefitfunction doesn't know anything about the
way layers are implemented, and can be used with other frameworks if desired. It
only relies on optax.TLDR;importjaximportjax.numpyasjnpimportjax.randomasrndimportoptaximportjynximportjynx.callbacksascbimportjynx.layersasnnfromjaximportArraydefdata_iter(key:Array):fromitertoolsimportrepeatx=jnp.linspace(-1,1,100).reshape(-1,1)y=jnp.cos(x)+0.05*rnd.normal(key,x.shape)returnrepeat((x,y))defloss_fn(net:nn.Module[[Array],Array],batch:tuple[Array,Array],# any type you wantkey:Array)->Array:x,y=batchy_pred=net(x,key=key)returnoptax.l2_loss(y_pred,y).mean()defmake_model(key:Array)->nn.Module[[Array],Array]:k1,k2,k3=rnd.split(key,3)net=nn.sequential(nn.linear(1,32,key=k1),jax.nn.relu,nn.linear(32,32,key=k2),jax.nn.relu,nn.linear(32,1,key=k3),)# or use:# net = nn.mlp(#     [1, 32, 32, 1],#     activation=jax.nn.silu,#     key=key,# )returnnetk1,k2,k3=rnd.split(rnd.PRNGKey(0),3)res=jynx.fit(make_model(k1),loss_fn=loss_fn,data_iter=data_iter(k2),optimizer=optax.adam(1e-3),key=k3,callbacks=[cb.ConsoleLogger(metrics=["loss"]),cb.TensorBoardLogger("tboard"),# requires tensorboardxcb.MlflowLogger(),# requires mlflowcb.CheckPoint("latest.pk"),# requires cloudpicklecb.EarlyStopping(monitor="loss",steps_without_improvement=500,),],)print("final loss",res.logs["loss"])net=res.paramsLayersCurrently implemented modules:SequentialParallelRecurrent: likeSequentialbut passes state, used for RNNsDenselyConnected: DenseNetLinearConvandConvTransposeEmbeddingFn: activation functionStarFn: equivalent offn(*x)Static: wraps an object to be ignored by jaxReshapeDropoutPooling layers:AvgPoolng,MaxPooling,MinPoolingNorm: layer norm, batch norm etc.SkipConnectionRNN layers:RNNCell,GRUCell,LSTMCellAttention,TransformerEncoderBlock,TransformerDecoderBlockConstructors for full networks:mlptransformer_encodertransformer_decoderrnn,lstm, andgruMore to come...How layers workLayers are simple pytree containers with a__call__method. To define new
modules easily, we provide a basePyTreeclass. Using this is not at all a
requirement, it just makes most definitions simpler. Layers that
don't require any static data can just as easily be defined asNamedTuples.classMyLinear(NamedTuple):weight:Arraybias:Arraydef__call__(self,x,*,key=None):returnx@self.weight+self.biasWe provide initialization with factory functions instead of in the
contructor. This makes flattening and unflattening pytrees much simpler. For
example:defmy_linear(size_in,size_out,*,key):w_init=jax.nn.initializers.kaiming_normal()returnMyLinear(w_init(key,(size_in,size_out)),jnp.zeros((size_out,)),)So layers are just 'dumb' containers.
ThePyTreebase class converts the inheriting class to a dataclass and
registers the type as ajaxpytreeclassMyDense(PyTree):weight:Arraybias:Arrayactivation:Callable[[Array],Array]=static(default=jax.nn.relu)def__call__(self,x,*,key=None):returnself.activation(x@self.weight+self.bias)ThefitfunctionTODOCallbacksTODO