Package: tensorwrap
TensorWrapTensorWrap - A full-fledged Deep Learning Library based on JAX and TensorFlow.|Install guideWhat is TensorWrap?TensorWrap is high performance neural network library that acts as a wrapper aroundJAX(another high performance machine learning library), bringing the familiar elements of theTensorFlow(2.x.x). This is currently aimed towards prototyping over deployment, in the current state.TensorWrap works by creating a layer of abstraction over JAX's low level api and introducing similar TensorFlow-like component's while supporting its own explicit and magic free design philosophy. This allows TensorWrap to be fast and efficient, while remaining nearly fully compatible with all custom operations and other tools from the JAX ecosystem. Additionally, this library adds additional features and leverages JAX's optimizations, making it more friendly towards research and educational audiences.This is a personal project, not professionally affliated with Google in any way. Expect bugs and several incompatibilities/difference from the original libraries.
Please help by trying it out,reporting
bugs, and letting me know what you
think!ContentsExamplesCurrent gimmicksInstallationNeural net librariesCitationsReference documentationExamplesCustom LayersimporttensorwrapastffromtensorwrapimportnnclassDense(nn.layers.Layer):def__init__(self,units)->None:super().__init__()# Needed for tracking trainable_variables.self.units=units# Defining the output shapedefbuild(self,input_shape:tuple)->None:super().build()# Required for letting model know that layer is built.input_shape=tf.shape(input_shape)# Getting appropriate input shape# Naming each parameter to later access from model.trainable_variablesself.kernel=self.add_weights([input_shape,self.units],initializer='glorot_uniform',name='kernel')self.bias=self.add_weights([self.units],initializer='zeros',name='bias')# Use call not __call__ to define the flow. To support JIT compilation, we use staticmethod.@staticmethod@tf.functiondefcall(params,inputs):returninputs@params['kernel']+params['bias']# Using params as an input, allows use to pass in the model.trainable_variables later.Just In Time Compiling with tf.functionimporttensorwrapastffromtensorwrapimportnntf.test.is_device_available(device_type='cuda')@tf.functiondefmse(y_pred,y_true):returntf.mean(tf.square(y_pred-y_true))print(mse(100,102))Custom ModelsimporttensorwrapastffromtensorwrapimportnnclassSequential(nn.Model):def__init__(self,layers:list)->None:super().__init__(name="Sequential")# Starts the tracking of internal variables. Allows for name definition.self.layers=layersdef__call__(self,inputs):x=inputsforlayerinself.layers:x=layer(x)returnxmodel=Sequential([nn.layers.Dense(100),nn.layers.Dense(10)])Current GimmicksCurrent models are all compiled by JAX's internal jit, so any error may remain a bit more cryptic than PyTorchs. However, this problem is still being worked on.Also, usingtensorwrap.Moduleis currently not recommended, since other superclasses offer more functionality and ease of use.Graph execution is currently not available, which means that all exported models can only be deployed within a python environment.InstallationThe device installation of TensorWrap depends on its backend, being JAX. Thus, our normal install will be covering only the cpu version. For gpu version, please checkJAX's documentation.pipinstall--upgradepip
pipinstall--upgradetensorwrapOn Linux, it is often necessary to first updatepipto a version that supportsmanylinux2014wheels. Also note that for Linux, we currently release wheels forx86_64architectures only, other architectures require building from source. Trying to pip install with other Linux architectures may lead tojaxlibnot being installed alongsidejax, althoughjaxmay successfully install (but fail at runtime).Thesepipinstallations do not work with Windows, and may fail silently; seeabove.NoteIf any problems occur with cuda installation, please visit theJAXgithub page, in order to understand the problem with lower API installation.CitationsThis project have been heavily inspired byTensorFlowand once again, is built on the open-source machine learning XLA frameworkJAX. Therefore, I recognize the authors of JAX and TensorFlow for the exceptional work they have done and understand that my library doesn't profit in any sort of way, since it is merely an add-on to the already existing community.@software{jax2018github,
  author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang},
  title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs},
  url = {http://github.com/google/jax},
  version = {0.3.13},
  year = {2018},
}Reference documentationFor details about the TensorWrap API, see the
[main documentation] (coming soon!)For details about JAX, see thereference documentation.For documentation on TensorFlow API, see theAPI documentation