Decoupling: Fast tensor decoupling in Jax.

Tensor decoupling is a methodology for decoupling multivariate functions using tensor decompositions.

This library aims to be the reference implementation of tensor decoupling algorithms.

Usage

import jax, jax.numpy as jnp

from decoupling import Algorithm
from decoupling import utils

def target(x): # define a simple polynomial
    return jnp.array([
        x[0]**3 + x[1]**2 + x[0]*x[1],
        x[1]**3 + x[0]**2 + x[0]*x[1],
    ])

key = jax.random.key(0)
rank = 4 

# collect information (inputs, outputs, jacobians)
info = utils.collect_information_from_random(target, N=100, key=key, ninputs=2)

decoupling = Algorithm(rank, key).run(*info) # compute decoupling

e1, e2 = utils.function_error_from_callable(target, decoupling, info[0])
print(f'Errors: e1={e1:.2f}, e2={e2:.2f}') # compare to target
# Errors: e1=1.20, e2=1.29 

Installation

You can easily get decoupling from PyPI:

pip install decoupling

Otherwise:

git clone git@github.com:mrochk/decoupling.git
pip install decoupling

Methodology

Tensor decoupling algorithms are used to find decoupled representations of target multivariate functions. This is illustrated below.

This decoupled representation is a one-hidden-layer MLP, meaning that tensor decoupling approaches can be used to compress or build neural networks.

You can read more about the basic methodology in this paper: https://arxiv.org/abs/1410.4060.

Our goal is to keep the source code as simple as possible, while being fast by leveraging Jax's JIT compiler.

Examples

To learn more about how this library works, feel free to take a look at some examples.

Testing

uv run python -m pytest -s # or ./test.sh

Context

The development of this library was started in the context of my master's thesis.

Citing Decoupling

@software{decoupling2026github,
  author = {Maxime Rochkoulets},
  title = {Decoupling: Fast tensor decoupling in {JAX}.},
  url = {http://github.com/mrochk/decoupling},
  version = {0.0.9},
  year = {2026},
}