pytreeclass 0.0.9

JAX compatible dataclass.

Write pytorch-like layers with rich visualizations in JAX.

Installation |Description |Quick Example |Filtering |StatefulComputation |Applications |More |Acknowledgements

🛠️ Installation

pip install pytreeclass

📖 Description

PyTreeClass offers a JAX compatible dataclass like datastructure with the following functionalities

• 🏗️ Create PyTorch like NN classes
• 🎨 Visualize for pytrees decorated with @pytc.treeclass.
• ☝️ Filtering by boolean masking similar to jax.numpy.at

⏩ Quick Example

🏗️ Create simple MLP

For Autoencoder example from scratch see here

import jax
from jax import numpy as jnp
import pytreeclass as pytc
import matplotlib.pyplot as plt

@pytc.treeclass
class Linear :
   # Any variable not wrapped with @pytc.treeclass
   # should be declared as a dataclass field here
   weight : jnp.ndarray
   bias   : jnp.ndarray

   def __init__(self,key,in_dim,out_dim):
       self.weight = jax.random.normal(key,shape=(in_dim, out_dim)) * jnp.sqrt(2/in_dim)
       self.bias = jnp.ones((1,out_dim))

   def __call__(self,x):
       return x @ self.weight + self.bias

@pytc.treeclass
class StackedLinear:

    def __init__(self,key,in_dim,out_dim,hidden_dim):
        keys= jax.random.split(key,3)

        # Declaring l1,l2,l3 as dataclass_fields is optional
        # as l1,l2,l3 are Linear class that is wrapped with @pytc.treeclass
        # To strictly include nodes defined in dataclass fields 
        # use `@pytc.treeclass(field_only=True)`
        self.l1 = Linear(key=keys[0],in_dim=in_dim,out_dim=hidden_dim)
        self.l2 = Linear(key=keys[1],in_dim=hidden_dim,out_dim=hidden_dim)
        self.l3 = Linear(key=keys[2],in_dim=hidden_dim,out_dim=out_dim)

    def __call__(self,x):
        x = self.l1(x)
        x = jax.nn.tanh(x)
        x = self.l2(x)
        x = jax.nn.tanh(x)
        x = self.l3(x)

        return x
        
>>> model = StackedLinear(in_dim=1,out_dim=1,hidden_dim=10,key=jax.random.PRNGKey(0))

>>> x = jnp.linspace(0,1,100)[:,None]
>>> y = x**3 + jax.random.uniform(jax.random.PRNGKey(0),(100,1))*0.01

🎨 Visualize

summary	tree_box	tree_diagram
>>> print(model.summary()) ┌──────┬───────┬───────┬─────────────────┐ │Type │Param #│Size │Config │ ├──────┼───────┼───────┼─────────────────┤ │Linear│20 │80.00B │weight=f32[1,10] │ │ │(0) │(0.00B)│bias=f32[1,10] │ ├──────┼───────┼───────┼─────────────────┤ │Linear│110 │440.00B│weight=f32[10,10]│ │ │(0) │(0.00B)│bias=f32[1,10] │ ├──────┼───────┼───────┼─────────────────┤ │Linear│11 │44.00B │weight=f32[10,1] │ │ │(0) │(0.00B)│bias=f32[1,1] │ └──────┴───────┴───────┴─────────────────┘ Total # : 141(0) Dynamic #: 141(0) Static/Frozen #: 0(0) ------------------------------------------ Total size : 564.00B(0.00B) Dynamic size: 564.00B(0.00B) Static/Frozen size: 0.00B(0.00B) ==========================================	>>> print(model.tree_box(array=x)) # using jax.eval_shape (no-flops operation) # ** note ** : the created modules # in __init__ should be in the same order # where they are called in __call__ ┌─────────────────────────────────────┐ │StackedLinear(Parent) │ ├─────────────────────────────────────┤ │┌────────────┬────────┬─────────────┐│ ││ │ Input │ f32[100,1] ││ ││ Linear(l1) │────────┼─────────────┤│ ││ │ Output │ f32[100,10] ││ │└────────────┴────────┴─────────────┘│ │┌────────────┬────────┬─────────────┐│ ││ │ Input │ f32[100,10] ││ ││ Linear(l2) │────────┼─────────────┤│ ││ │ Output │ f32[100,10] ││ │└────────────┴────────┴─────────────┘│ │┌────────────┬────────┬─────────────┐│ ││ │ Input │ f32[100,10] ││ ││ Linear(l3) │────────┼─────────────┤│ ││ │ Output │ f32[100,1] ││ │└────────────┴────────┴─────────────┘│ └─────────────────────────────────────┘	>>> print(model.tree_diagram()) StackedLinear ├── l1=Linear │ ├── weight=f32[1,10] │ └── bias=f32[1,10] ├── l2=Linear │ ├── weight=f32[10,10] │ └── bias=f32[1,10] └──l3=Linear ├── weight=f32[10,1] └── bias=f32[1,1]

mermaid.io (Native support in Github/Notion)

# generate mermaid diagrams # print(pytc.tree_viz.tree_mermaid(model)) # generate core syntax >>> pytc.tree_viz.save_viz(model,filename="test_mermaid",method="tree_mermaid_md") # use `method="tree_mermaid_html"` to save as html flowchart LR id15696277213149321320[StackedLinear] id15696277213149321320 --> id159132120600507116(l1\nLinear) id159132120600507116 --- id7500441386962467209["weight\nf32[1,10]"] id159132120600507116 --- id10793958738030044218["bias\nf32[1,10]"] id15696277213149321320 --> id10009280772564895168(l2\nLinear) id10009280772564895168 --- id11951215191344350637["weight\nf32[10,10]"] id10009280772564895168 --- id1196345851686744158["bias\nf32[1,10]"] id15696277213149321320 --> id7572222925824649475(l3\nLinear) id7572222925824649475 --- id4749243995442935477["weight\nf32[10,1]"] id7572222925824649475 --- id8042761346510512486["bias\nf32[1,1]"] ✨ Generate shareable vizualization links ✨ >>> pytc.tree_viz.tree_mermaid(model,link=True) 'Open URL in browser: https://pytreeclass.herokuapp.com/temp/?id=*********'

✂️ Model surgery

# freeze l1
>>> model.l1 = model.l1.freeze()

# set negative values in l2 to 0
>>> model.l2 = model.l2.at[model.l2<0].set(0)

# apply sin(x) to all values in l3
>>> model.l3 = model.l3.at[...].apply(jnp.sin)

# frozen nodes are marked with #
>>> print(model.tree_diagram())
StackedLinear
    ├── l1=Linear
    │   ├#─ weight=f32[1,10]
    │   └#─ bias=f32[1,10]  
    ├── l2=Linear
    │   ├── weight=f32[10,10]
    │   └── bias=f32[1,10]  
    └── l3=Linear
        ├── weight=f32[10,1]
        └── bias=f32[1,1] 

☝️ Filtering with .at[]

PyTreeClass offers four means of filtering:

1. Filter by value
2. Filter by field name
3. Filter by field type
4. Filter by field metadata.

The following example demonstrates the usage the filtering. Suppose you have the following (Multilayer perceptron) MLP class

• Note in StackedLinear l1 and l2 has a description in field metadata.

Model definition

import jax
from jax import numpy as jnp
import pytreeclass as pytc
import matplotlib.pyplot as plt
from dataclasses import  field 

@pytc.treeclass
class Linear :
   weight : jnp.ndarray
   bias   : jnp.ndarray

   def __init__(self,key,in_dim,out_dim):
       self.weight = jax.random.normal(key,shape=(in_dim, out_dim)) * jnp.sqrt(2/in_dim)
       self.bias = jnp.ones((1,out_dim))

   def __call__(self,x):
       return x @ self.weight + self.bias

@pytc.treeclass
class StackedLinear:
    l1 : Linear = field(metadata={"description": "First layer"})
    l2 : Linear = field(metadata={"description": "Second layer"})

    def __init__(self,key,in_dim,out_dim,hidden_dim):
        keys= jax.random.split(key,3)

        self.l1 = Linear(key=keys[0],in_dim=in_dim,out_dim=hidden_dim)
        self.l2 = Linear(key=keys[2],in_dim=hidden_dim,out_dim=out_dim)

    def __call__(self,x):
        x = self.l1(x)
        x = jax.nn.tanh(x)
        x = self.l2(x)

        return x
        
model = StackedLinear(in_dim=1,out_dim=1,hidden_dim=5,key=jax.random.PRNGKey(0))

• Raw model values before any filtering.

>>> print(model)

StackedLinear(
  l1=Linear(
    weight=[[-1.6248673  -2.8383057   1.3969219   1.3169124  -0.40784812]],
    bias=[[1. 1. 1. 1. 1.]]),
  l2=Linear(
    weight=
      [[ 0.98507565]
       [ 0.99815285]
       [-1.0687716 ]
       [-0.19255024]
       [-1.2108876 ]],
    bias=[[1.]]))

Filter by value

• Get all negative values

>>> print(model.at[model<0].get())

StackedLinear(
  l1=Linear(
    weight=[-1.6248673  -2.8383057  -0.40784812],
    bias=[]),
  l2=Linear(
    weight=[-1.0687716  -0.19255024 -1.2108876 ],
    bias=[]))

• Set negative values to 0

>>> print(model.at[model<0].set(0))

StackedLinear(
  l1=Linear(
    weight=[[0.        0.        1.3969219 1.3169124 0.       ]],
    bias=[[1. 1. 1. 1. 1.]]),
  l2=Linear(
    weight=
      [[0.98507565]
       [0.99815285]
       [0.        ]
       [0.        ]
       [0.        ]],
    bias=[[1.]]))

• Apply f(x)=x^2 to negative values

>>> print(model.at[model<0].apply(lambda x:x**2))

StackedLinear(
  l1=Linear(
    weight=[[2.6401937  8.05598    1.3969219  1.3169124  0.16634008]],
    bias=[[1. 1. 1. 1. 1.]]),
  l2=Linear(
    weight=
      [[0.98507565]
       [0.99815285]
       [1.1422727 ]
       [0.03707559]
       [1.4662486 ]],
    bias=[[1.]]))

• Sum all negative values

>>> print(model.at[model<0].reduce_sum())
-7.3432307

Filter by field name

• Get all fields named l1

>>> print(model.at[model == "l1"].get())

StackedLinear(
  l1=Linear(
    weight=[-1.6248673  -2.8383057   1.3969219   1.3169124  -0.40784812],
    bias=[1. 1. 1. 1. 1.]),
  l2=Linear(weight=[],bias=[]))

Filter by field type

• Get all fields of Linear type

>>> print(model.at[model == Linear].get())

StackedLinear(
  l1=Linear(
    weight=[-1.6248673  -2.8383057   1.3969219   1.3169124  -0.40784812],
    bias=[1. 1. 1. 1. 1.]),
  l2=Linear(
    weight=[ 0.98507565  0.99815285 -1.0687716  -0.19255024 -1.2108876 ],
    bias=[1.]))

Filter by field metadata

• Get all fields of with {"description": "First layer"} in their metadata

>>> print(model.at[model == {"description": "First layer"}].get())

StackedLinear(
  l1=Linear(
    weight=[-1.6248673  -2.8383057   1.3969219   1.3169124  -0.40784812],
    bias=[1. 1. 1. 1. 1.]),
  l2=Linear(weight=[],bias=[]))

🤯 Application : Filtering PyTrees by boolean masking

• Manipulate certain modules attributes values.
• Set certain modules (e.g. Dropout) to eval mode

• Model definition

import jax
from jax import numpy as jnp
import jax.random as jr
import pytreeclass as pytc

@pytc.treeclass
class Linear :
   weight : jnp.ndarray
   bias   : jnp.ndarray

   def __init__(self,key,in_dim,out_dim):
       self.weight = jax.random.normal(key,shape=(in_dim, out_dim)) * jnp.sqrt(2/in_dim)
       self.bias = jnp.ones((1,out_dim))

   def __call__(self,x):
       return x @ self.weight + self.bias


@pytc.treeclass
class Dropout:
    p: float
    eval : bool | None

    def __init__(self, p: float = 0.5, eval: bool | None = None):
        """p : probability of an element to be zeroed out"""
        self.p = p
        self.eval = eval

    def __call__(self, x, *, key=jr.PRNGKey(0)):
        return ( 
            x if (self.eval is True) 
            else 
            jnp.where(jr.bernoulli(key, (1 - self.p), x.shape), x / (1 - self.p), 0)
        )

@pytc.treeclass
class LinearWithDropout:
    def __init__(self):
        self.l1 = Linear(key=jr.PRNGKey(0), in_dim=1, out_dim=5)
        self.d1 = Dropout(p = 1.) # zero out all elements 

    def __call__(self, x):
        x = self.l1(x)
        x = self.d1(x)
        return x

Linear module with full dropout

>>> model = LinearWithDropout()
>>> print(model(jnp.ones((1,1))))

[[0. 0. 0. 0. 0.]]

Disable Dropout

• using boolean masking with .at[].set() to disable Dropout

>>> mask = (model == "eval")
>>> model_no_dropout = model.at[mask].set(True, is_leaf = lambda x:x is None)
>>> print(model_no_dropout(jnp.ones((1,1))))

[[ 1.2656513  -0.8149204   0.61661845  2.7664368   1.3457328 ]]

Set Linear module bias to 0

• Combining attribute name mask and class type mask

>>> mask = (model == "bias") & (model == Linear) 
>>> model_no_linear_bias = model.at[mask ].set(0)
>>> print(model_no_linear_bias)

LinearWithDropout(
 l1=Linear(
   weight=[[ 0.26565132 -1.8149204  -0.38338155  1.7664368   0.34573284]],
   bias=[[0. 0. 0. 0. 0.]]),
 d1=Dropout(p=1.0,eval=None))

📜 Stateful computations

JAX reference Under jax.jit jax requires states to be explicit, this means that for any class instance; variables needs to be separated from the class and be passed explictly. However when using @pytc.treeclass no need to separate the instance variables ; instead the whole instance is passed as a state.

Using the following pattern,Updating state can be achieved under jax.jit

@pytc.treeclass
class Counter:
    calls : int = 0
    
    def increment(self):
        self.calls += 1 
        

>>> c = Counter()

@jax.jit
def update(c):
    c.increment()
    return c 

for i in range(10):
    c = update(c)

>>> print(c.calls)
10

The following code snippets compares between the two concepts by comparing MLP's implementation.

Explicit state

Class instance as state

import jax.numpy as jnp import jax.random as jr from jax.nn.initializers import he_normal from jax.tree_util import tree_map from jax import nn, value_and_grad,jit import pytreeclass as pytc def init_params(layers): keys = jr.split( jr.PRNGKey(0),len(layers)-1 ) params = list() init_func = he_normal() for key,n_in,n_out in zip( keys, layers[:-1], layers[1:] ): W = init_func(key,(n_in,n_out)) B = jr.uniform(key,shape=(n_out,)) params.append({'W':W,'B':B}) return params def fwd(params,x): *hidden,last = params for layer in hidden : x = nn.tanh(x@layer['W']+layer['B']) return x@last['W'] + last['B'] @value_and_grad def loss_func(params,x,y): pred = fwd(params,x) return jnp.mean((pred-y)**2) @jit def update(params,x,y): # gradient w.r.t to params value,grads= loss_func(params,x,y) params = tree_map( lambda x,y : x-1e-3*y, params,grads ) return value,params x = jnp.linspace(0,1,100).reshape(100,1) y = x**2 -1 params = init_params([1] +[5]*4+[1] ) epochs = 10_000 for _ in range(1,epochs+1): value , params = update(params,x,y) # print loss and epoch info if _ %(1_000) ==0: print(f'Epoch={_}\tloss={value:.3e}')

import jax.numpy as jnp import jax.random as jr from jax.nn.initializers import he_normal from jax.tree_util import tree_map from jax import nn, value_and_grad,jit import pytreeclass as pytc @pytc.treeclass class MLP: Layers : list def __init__(self,layers): keys = jr.split( jr.PRNGKey(0),len(layers)-1 ) self.Layers = list() init_func = he_normal() for key,n_in,n_out in zip( keys, layers[:-1], layers[1:] ): W = init_func(key,(n_in,n_out)) B = jr.uniform(key,shape=(n_out,)) self.Layers.append({'W':W,'B':B}) def __call__(self,x): *hidden,last = self.Layers for layer in hidden : x = nn.tanh(x@layer['W']+layer['B']) return x@last['W'] + last['B'] @value_and_grad def loss_func(model,x,y): pred = model(x) return jnp.mean((pred-y)**2) @jit def update(model,x,y): # gradient w.r.t to model value , grads= loss_func(model,x,y) model = tree_map( lambda x,y : x-1e-3*y, model,grads ) return value , model x = jnp.linspace(0,1,100).reshape(100,1) y = x**2 -1 model = MLP([1] +[5]*4+[1] ) epochs = 10_000 for _ in range(1,epochs+1): value , model = update(model,x,y) # print loss and epoch info if _ %(1_000) ==0: print(f'Epoch={_}\tloss={value:.3e}')

📝 Applications

• Physics informed neural network (PINN)

🔢 More

More compact boilerplate

Standard definition of nodes in __init__ and calling in __call__

@pytc.treeclass
class StackedLinear:
    def __init__(self,key,in_dim,out_dim,hidden_dim):
        keys= jax.random.split(key,3)
        self.l1 = Linear(key=keys[0],in_dim=in_dim,out_dim=hidden_dim)
        self.l2 = Linear(key=keys[1],in_dim=hidden_dim,out_dim=hidden_dim)
        self.l3 = Linear(key=keys[2],in_dim=hidden_dim,out_dim=out_dim)

    def __call__(self,x):
        x = self.l1(x)
        x = jax.nn.tanh(x)
        x = self.l2(x)
        x = jax.nn.tanh(x)
        x = self.l3(x)
        return x

Using register_node:

• More compact definition with node definition at runtime call
• The Linear layers are defined on the first call and retrieved on the subsequent calls
• This pattern is useful if module definition depends runtime data.

@pytc.treeclass
class StackedLinear:
    def __init__(self,key):
        self.keys = jax.random.split(key,3)

    def __call__(self,x):
        x = self.register_node(Linear(self.keys[0],x.shape[-1],10),name="l1")(x)
        x = jax.nn.tanh(x)
        x = self.register_node(Linear(self.keys[1],10,10),name="l2")(x)
        x = jax.nn.tanh(x)
        x = self.register_node(Linear(self.keys[2],10,x.shape[-1]),name="l3")(x)
        return x

Simple AutoEncoder from scratch

While jax.lax can be used to construct Convolution, Upsample, Maxpooling functions, in this example kernex is used for its clear syntax.

AE Construction

from typing import Sequence

import jax
import jax.numpy as jnp
import jax.random as jr
import pytreeclass as pytc  # dataclass-like decorator for JAX

import kernex as kex # for stencil computations


@pytc.treeclass
class Conv2D:

    weight: jnp.ndarray
    bias: jnp.ndarray

    # define these variabels here
    # to be used in __call__
    in_channels: int = pytc.static_field()
    out_channels: int = pytc.static_field()
    kernel_size: Sequence[int] = pytc.static_field()
    padding: Sequence[str] = pytc.static_field()
    strides: Sequence[int] = pytc.static_field()

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        strides=1,
        padding=("same", "same"),
        key=jax.random.PRNGKey(0),
        kernel_initializer=jax.nn.initializers.kaiming_uniform(),
    ):

        self.weight = kernel_initializer(key, (out_channels, in_channels, *kernel_size))
        self.bias = jnp.zeros((out_channels, *((1,) * len(kernel_size))))

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.strides = strides
        self.padding = ("valid",) + padding

    def __call__(self, x):
        @kex.kmap(
            kernel_size=(self.in_channels, *self.kernel_size),
            strides=self.strides,
            padding=self.padding,
        )
        def _conv2d(x, w):
            return jnp.sum(x * w)

        @jax.vmap  # vectorize on batch dimension
        def fwd_image(image):
            # filters shape is OIHW
            # vectorize on filters output dimension
            return jax.vmap(lambda w: _conv2d(image, w))(self.weight)[:, 0] + (
                self.bias if self.bias is not None else 0
            )

        return fwd_image(x)


@pytc.treeclass
class Upsample2D:
    scale_factor: int = pytc.static_field()

    def __call__(self, x):

        batch, channel, row, col = x.shape

        @kex.kmap(
            kernel_size=(channel, row, col),
            strides=(1, 1, 1),
            padding="valid",
            relative=False,
        )
        def __upsample2D(x):
            return x.repeat(self.scale_factor, axis=2).repeat(self.scale_factor, axis=1)

        def _upsample2D(batch):
            return jnp.squeeze(
                jax.vmap(__upsample2D, in_axes=(0,))(batch), axis=tuple(range(1, 4))
            )

        return _upsample2D(x)


@pytc.treeclass
class MaxPool2D:

    kernel_size: tuple[int, int] = pytc.static_field(default=(2, 2))
    strides: int = pytc.static_field(default=2)
    padding: str | int = pytc.static_field(default="valid")

    def __call__(self, x):
        @jax.vmap  # apply on batch dimension
        @jax.vmap  # apply on channels dimension
        @kex.kmap(
            kernel_size=self.kernel_size, strides=self.strides, padding=self.padding
        )
        def _maxpool2d(x):
            return jnp.max(x)

        return _maxpool2d(x)


@pytc.treeclass
class AutoEncoder:
    def __init__(self, in_channels, out_channels, key):
        keys = jr.split(key, 5)

        self.l1 = MaxPool2D()
        self.l2 = Conv2D(in_channels, 16, (3, 3), key=keys[0])

        self.l3 = MaxPool2D()
        self.l4 = Conv2D(16, 32, (3, 3), key=keys[1])

        self.l5 = Upsample2D(scale_factor=2)
        self.l6 = Conv2D(32, 16, (3, 3), key=keys[2])
        
        self.l7 = Upsample2D(scale_factor=2)
        self.l8 = Conv2D(16, 1, (3, 3), key=keys[3])

        self.l9 = Conv2D(1, out_channels, (1, 1), key=keys[4])

    def __call__(self, x):
        x = self.l1(x)
        x = self.l2(x)
        x = jax.nn.relu(x)

        x = self.l3(x)
        x = self.l4(x)
        x = jax.nn.relu(x)

        x = self.l5(x)
        x = self.l6(x)
        x = jax.nn.relu(x)

        x = self.l7(x)
        x = self.l8(x)
        x = jax.nn.relu(x)

        x = self.l9(x)

        return x


ae = AutoEncoder(1, 1, jax.random.PRNGKey(0))

Model summary

┌──────────┬───────┬───────┬─────────────────────┐
│Type      │Param #│Size   │Config               │
├──────────┼───────┼───────┼─────────────────────┤
│MaxPool2D │0      │0.00B  │                     │
│          │(0)    │(0.00B)│                     │
├──────────┼───────┼───────┼─────────────────────┤
│Conv2D    │160    │640.00B│weight=f32[16,1,3,3] │
│          │(0)    │(0.00B)│bias=f32[16,1,1]     │
├──────────┼───────┼───────┼─────────────────────┤
│MaxPool2D │0      │0.00B  │                     │
│          │(0)    │(0.00B)│                     │
├──────────┼───────┼───────┼─────────────────────┤
│Conv2D    │4,640  │18.12KB│weight=f32[32,16,3,3]│
│          │(0)    │(0.00B)│bias=f32[32,1,1]     │
├──────────┼───────┼───────┼─────────────────────┤
│Upsample2D│0      │0.00B  │                     │
│          │(0)    │(0.00B)│                     │
├──────────┼───────┼───────┼─────────────────────┤
│Conv2D    │4,624  │18.06KB│weight=f32[16,32,3,3]│
│          │(0)    │(0.00B)│bias=f32[16,1,1]     │
├──────────┼───────┼───────┼─────────────────────┤
│Upsample2D│0      │0.00B  │                     │
│          │(0)    │(0.00B)│                     │
├──────────┼───────┼───────┼─────────────────────┤
│Conv2D    │145    │580.00B│weight=f32[1,16,3,3] │
│          │(0)    │(0.00B)│bias=f32[1,1,1]      │
├──────────┼───────┼───────┼─────────────────────┤
│Conv2D    │2      │8.00B  │weight=f32[1,1,1,1]  │
│          │(0)    │(0.00B)│bias=f32[1,1,1]      │
└──────────┴───────┴───────┴─────────────────────┘
Total # :		9,571(0)
Dynamic #:		9,571(0)
Static/Frozen #:	0(0)
--------------------------------------------------
Total size :		37.39KB(0.00B)
Dynamic size:		37.39KB(0.00B)
Static/Frozen size:	0.00B(0.00B)
==================================================

Model diagram

Note : static_field(untrainable) is marked with x

AutoEncoder
    ├── l1=MaxPool2D
    │   ├x─ kernel_size=(2, 2)
    │   ├x─ strides=2
    │   └x─ padding='valid' 
    ├── l2=Conv2D
    │   ├── weight=f32[16,1,3,3]
    │   ├── bias=f32[16,1,1]
    │   ├x─ in_channels=1
    │   ├x─ out_channels=16
    │   ├x─ kernel_size=(3, 3)
    │   ├x─ padding=('valid', 'same', 'same')
    │   └x─ strides=1   
    ├── l3=MaxPool2D
    │   ├x─ kernel_size=(2, 2)
    │   ├x─ strides=2
    │   └x─ padding='valid' 
    ├── l4=Conv2D
    │   ├── weight=f32[32,16,3,3]
    │   ├── bias=f32[32,1,1]
    │   ├x─ in_channels=16
    │   ├x─ out_channels=32
    │   ├x─ kernel_size=(3, 3)
    │   ├x─ padding=('valid', 'same', 'same')
    │   └x─ strides=1   
    ├── l5=Upsample2D
    │   └x─ scale_factor=2  
    ├── l6=Conv2D
    │   ├── weight=f32[16,32,3,3]
    │   ├── bias=f32[16,1,1]
    │   ├x─ in_channels=32
    │   ├x─ out_channels=16
    │   ├x─ kernel_size=(3, 3)
    │   ├x─ padding=('valid', 'same', 'same')
    │   └x─ strides=1   
    ├── l7=Upsample2D
    │   └x─ scale_factor=2  
    ├── l8=Conv2D
    │   ├── weight=f32[1,16,3,3]
    │   ├── bias=f32[1,1,1]
    │   ├x─ in_channels=16
    │   ├x─ out_channels=1
    │   ├x─ kernel_size=(3, 3)
    │   ├x─ padding=('valid', 'same', 'same')
    │   └x─ strides=1   
    └── l9=Conv2D
        ├── weight=f32[1,1,1,1]
        ├── bias=f32[1,1,1]
        ├x─ in_channels=1
        ├x─ out_channels=1
        ├x─ kernel_size=(1, 1)
        ├x─ padding=('valid', 'same', 'same')
        └x─ strides=1                               

Shape propagration

>>> x = jnp.ones([1, 1, 100, 100])
>>> print(ae.tree_box(array=x))
┌───────────────────────────────────────────────┐
│AutoEncoder(Parent)                            │
├───────────────────────────────────────────────┤
│┌───────────────┬────────┬──────────────────┐  │
││               │ Input  │ f32[1,1,100,100] │  │
││ MaxPool2D(l1) │────────┼──────────────────┤  │
││               │ Output │ f32[1,1,50,50]   │  │
│└───────────────┴────────┴──────────────────┘  │
│┌────────────┬────────┬─────────────────┐      │
││            │ Input  │ f32[1,1,50,50]  │      │
││ Conv2D(l2) │────────┼─────────────────┤      │
││            │ Output │ f32[1,16,50,50] │      │
│└────────────┴────────┴─────────────────┘      │
│┌───────────────┬────────┬─────────────────┐   │
││               │ Input  │ f32[1,16,50,50] │   │
││ MaxPool2D(l3) │────────┼─────────────────┤   │
││               │ Output │ f32[1,16,25,25] │   │
│└───────────────┴────────┴─────────────────┘   │
│┌────────────┬────────┬─────────────────┐      │
││            │ Input  │ f32[1,16,25,25] │      │
││ Conv2D(l4) │────────┼─────────────────┤      │
││            │ Output │ f32[1,32,25,25] │      │
│└────────────┴────────┴─────────────────┘      │
│┌────────────────┬────────┬─────────────────┐  │
││                │ Input  │ f32[1,32,25,25] │  │
││ Upsample2D(l5) │────────┼─────────────────┤  │
││                │ Output │ f32[1,32,50,50] │  │
│└────────────────┴────────┴─────────────────┘  │
│┌────────────┬────────┬─────────────────┐      │
││            │ Input  │ f32[1,32,50,50] │      │
││ Conv2D(l6) │────────┼─────────────────┤      │
││            │ Output │ f32[1,16,50,50] │      │
│└────────────┴────────┴─────────────────┘      │
│┌────────────────┬────────┬───────────────────┐│
││                │ Input  │ f32[1,16,50,50]   ││
││ Upsample2D(l7) │────────┼───────────────────┤│
││                │ Output │ f32[1,16,100,100] ││
│└────────────────┴────────┴───────────────────┘│
│┌────────────┬────────┬───────────────────┐    │
││            │ Input  │ f32[1,16,100,100] │    │
││ Conv2D(l8) │────────┼───────────────────┤    │
││            │ Output │ f32[1,1,100,100]  │    │
│└────────────┴────────┴───────────────────┘    │
│┌────────────┬────────┬──────────────────┐     │
││            │ Input  │ f32[1,1,100,100] │     │
││ Conv2D(l9) │────────┼──────────────────┤     │
││            │ Output │ f32[1,1,100,100] │     │
│└────────────┴────────┴──────────────────┘     │
└───────────────────────────────────────────────┘

Mermaid diagram

flowchart LR
    id15696277213149321320[AutoEncoder]
    id15696277213149321320 --> id159132120600507116(l1\nMaxPool2D)
    id159132120600507116 --x id7500441386962467209["kernel_size\n(2, 2)"]
    id159132120600507116 --x id10793958738030044218["strides\n2"]
    id159132120600507116 --x id16245750007010064142["padding\n'valid'"]
    id15696277213149321320 --> id10009280772564895168(l2\nConv2D)
    id10009280772564895168 --- id11951215191344350637["weight\nf32[16,1,3,3]"]
    id10009280772564895168 --- id1196345851686744158["bias\nf32[16,1,1]"]
    id10009280772564895168 --x id6648137120666764082["in_channels\n1"]
    id10009280772564895168 --x id8609436656910886517["out_channels\n16"]
    id10009280772564895168 --x id14061227925890906441["kernel_size\n(3, 3)"]
    id10009280772564895168 --x id16022527462135028876["padding\n('valid', 'same', 'same')"]
    id10009280772564895168 --x id869300739493054269["strides\n1"]
    id15696277213149321320 --> id7572222925824649475(l3\nMaxPool2D)
    id7572222925824649475 --x id4749243995442935477["kernel_size\n(2, 2)"]
    id7572222925824649475 --x id8042761346510512486["strides\n2"]
    id7572222925824649475 --x id17892909998474900538["padding\n'valid'"]
    id15696277213149321320 --> id10865740276892226484(l4\nConv2D)
    id10865740276892226484 --- id7858522665561710831["weight\nf32[32,16,3,3]"]
    id10865740276892226484 --- id11152040016629287840["bias\nf32[32,1,1]"]
    id10865740276892226484 --x id2555444594884124276["in_channels\n16"]
    id10865740276892226484 --x id118386748143878583["out_channels\n32"]
    id10865740276892226484 --x id9968535400108266635["kernel_size\n(3, 3)"]
    id10865740276892226484 --x id7531477553368020942["padding\n('valid', 'same', 'same')"]
    id10865740276892226484 --x id10824994904435597951["strides\n1"]
    id15696277213149321320 --> id2269144855147062920(l5\nUpsample2D)
    id2269144855147062920 --x id599357636669938791["scale_factor\n2"]
    id15696277213149321320 --> id18278831082116368843(l6\nConv2D)
    id18278831082116368843 --- id5107325274042179099["weight\nf32[16,32,3,3]"]
    id18278831082116368843 --- id8400842625109756108["bias\nf32[16,1,1]"]
    id18278831082116368843 --x id18250991277074144160["in_channels\n32"]
    id18278831082116368843 --x id1765546739608714979["out_channels\n16"]
    id18278831082116368843 --x id7217338008588734903["kernel_size\n(3, 3)"]
    id18278831082116368843 --x id9178637544832857338["padding\n('valid', 'same', 'same')"]
    id18278831082116368843 --x id12472154895900434347["strides\n1"]
    id15696277213149321320 --> id9682235660371205279(l7\nUpsample2D)
    id9682235660371205279 --x id13157878626227910245["scale_factor\n2"]
    id15696277213149321320 --> id12975753011438782288(l8\nConv2D)
    id12975753011438782288 --- id16267157296346685599["weight\nf32[1,16,3,3]"]
    id12975753011438782288 --- id1113930573704710992["bias\nf32[1,1,1]"]
    id12975753011438782288 --x id10964079225669099044["in_channels\n16"]
    id12975753011438782288 --x id12925378761913221479["out_channels\n1"]
    id12975753011438782288 --x id11331140742258384578["kernel_size\n(3, 3)"]
    id12975753011438782288 --x id1891725493427812222["padding\n('valid', 'same', 'same')"]
    id12975753011438782288 --x id5185242844495389231["strides\n1"]
    id15696277213149321320 --> id10538695164698536595(l9\nConv2D)
    id10538695164698536595 --- id9065186100445270439["weight\nf32[1,1,1,1]"]
    id10538695164698536595 --- id12358703451512847448["bias\nf32[1,1,1]"]
    id10538695164698536595 --x id3762108029767683884["in_channels\n1"]
    id10538695164698536595 --x id1325050183027438191["out_channels\n1"]
    id10538695164698536595 --x id11175198834991826243["kernel_size\n(1, 1)"]
    id10538695164698536595 --x id8738140988251580550["padding\n('valid', 'same', 'same')"]
    id10538695164698536595 --x id12031658339319157559["strides\n1"]
   

📙 Acknowledgements

• Farid Talibli (for visualization link generation backend)
• Equinox
• Treex
• tree-math