Adding or changing Objax modules

This guide explains how to add a new module to Objax or change existing one.

In addition to this guide, consider looking at an example pull request which adds new module with documentation.

Writing code

When adding new module or function, you have to decide where to put it. Typical locations of new modules and functions are the following:

  • objax/functional contains various stateless functions (non-modules) which are used in machine learning. For example: loss functions, activations, stateless ops.

  • objax/io contains routines for model saving and loading.

  • objax/nn contains layers, which serve as building blocks for neural network. It also contains initializes for layer parameters.

  • objax/optimizer contains optimizers.

  • objax/privacy contains code for privacy-preserving training of neural networks.

  • objax/random contains routines for random number generation.

  • objax/zoo is a “model zoo” of various well-known neural network architectures.

When writing code we follow PEP8 style guide with the following two exceptions:

  • We maximum line length to 120 characters.

  • We allow to assign lambda, e.g. f = lambda x: x

  • In addition we recommend trying to keep APIs ordered alphabetically when feasible within source code files.

Note: Remember to add your new APIs in __all__ variable (also ordered alphabetically) at the top of the file to make them visible.

Script ./tests/ automatically checks majority of code style violations. PyCharm code formatter could be used to automatically reformat code.

Writing unit tests

Unit tests are required for most code changes and new code of Objax library. However tests are not required for examples.

All unit tests are placed into tests directory. They are grouped into different files based on what they are testing. For example tests/ contains unit tests for convolution modules.

We use Python unittest module for tests.

Writing documentation

Documentation for specific APIs is written inside docstrings within the code (example for Conv2D).

Other documentation is stored in docs subdirectory of Objax repository. It uses reST as a markup language, and Sphinx automatically generates documentation for


All public facing classes, functions and class methods should have a short docstring describing what they are doing. Functions and methods should also have a description of their arguments and return value.

To keep code easy to read, we recommend to write short and concise docstrings:

  • Try to keep description of classes and functions with 1-5 lines.

  • Try to fit description of each argument of each function within 1-2 lines.

  • Avoid putting examples or long descriptions into docstrings, those should go into reST docs.

Here is an example of how to write docstring for a function:

def cross_entropy_logits(logits: JaxArray, labels: JaxArray) -> JaxArray:
    """Computes the cross-entropy loss.

        logits: (batch, #class) tensor of logits.
        labels: (batch, #class) tensor of label probabilities (e.g. labels.sum(axis=1) must be 1)

        (batch,) tensor of the cross-entropies for each entry.
    return logsumexp(logits, axis=1) - (logits * labels).sum(1)

If you are only updating existing docstrings these changes will be automatically reflected in after pull request is merged into repository. When adding docstrings for new classes and functions, you also may need to update reST files as described below.

reST documentation

Updates of reST files are required either when new APIs are added (new function, new module) or when other (non API) documentation is needed.

Most of the API documentation is located in docs/source/objax They are grouped into different .rst files by the name of python package, for example docs/source/objax/nn.rst contains documentation for objax.nn package.

To add new class or function, you typically need to add name of the class or function into autosummary section and add autoclass or autofunction section for new class/function. Here is an example of changes which are needed to add documentation for Conv2D module:

.. autosummary::



.. autoclass:: Conv2D

    Additional documentation (non-docstrings) for Conv2D goes here.

For reference about reST syntax, refer to reST documentation or cheat sheet.