Organising code in the pure / functional world

[mchagneux]

Hi all,

I often find it convenient to define different levels of functionality in terms of inheritence and to override some class methods depending on the use case. In practise, in my current project I have some loss function F(params) which needs to be computed differently depending on the nature of "params", the latter describing parameters of another parameter which can be either linear or nonlinear. Many subroutines used to compute F are common between the linear and nonlinear case, but some are not, and when I was on Pytorch I used to heavily rely on class interfaces and redefine only the methods that differed in child classes, ie. I would have had some BaseF class with common subroutines, then LinearF(BaseF) and NonLinearF(BaseF) that only redefined what differs.

This whole thing doesn't seem to fit well withing Jax philosophy. On this page https://jax.readthedocs.io/en/latest/jax-101/07-state.html I read at the end:

In our case, the CounterV2 class is nothing more than a namespace bringing all the functions that use CounterState into one location. Exercise for the reader: do you think it makes sense to keep it as a class?

The following alternatives seem a bit verbose to me:

• writing as many functions (with different names) for each use case and using control flow to select the correct ones.
• creating separate modules, one for each use case, with common functions in another (base) module. That's like reproducing the inheritance practicality with modules.

Is there another standard alternative that I'm not seeing ? One thing I have in mind is using inheritence but with only static methods (an no abstract methods). But this seems like the most degenarate case for using classes.

Thanks in advance!

[mchagneux]

I think it works with the classmethod! I thought I was gonna run into trouble with the implicit first argument "cls" when using jax functional transforms, but apparently that's handled appropriately. Still, if there's another, more standard way of proceeding, I'd be glad to know about it.

[YouJiacheng]

I find that jit does not support staticmethod.
But class attribute is ok:

class parent:
    foo = jax.jit(lambda x: x + 1)

class childA(parent):
    bar = jax.jit(lambda x: x + 2)

class childB(parent):
    foo = jax.jit(lambda x: x + 3)

print(parent.foo(0)) # 1
print(childA.foo(0)) # 1
print(childA.bar(0)) # 2
print(childB.foo(0)) # 3

[mchagneux]

Hi and thanks for your reply!

I've investigated the question a little bit and made some tests. To me classmethod is the best option, because staticmethod will not let me resolve the instantiated class automatically, which is convenient if I want to call the child (redefined) functions from a common function.

The following code runs without any error:

class F(metaclass=ABCMeta):

    @classmethod
    def compute(cls, x):
        return cls.specific_subroutine(x) + cls.common_subroutine(x)

    @classmethod
    def common_subroutine(cls, x):
        return x + 1 

    @classmethod
    @abstractmethod
    def specific_subroutine(cls, x):
        pass 
        
class LinearF(F):

    @classmethod
    def specific_subroutine(cls, x):
        return 2*x 

class NonLinearF(F):

    @classmethod 
    def specific_subroutine(cls, x):
        return x**2

def scanner(carry, x):
    return carry, LinearF.compute(x)


@jax.jit
def scan(xs): 
    _ , results = jax.lax.scan(scanner, init=None, xs = xs)
    return results 

print(scan(2*jax.numpy.ones((10,))))

So basically I can jit (or grad) anything here without any problem.

If I really need to jit specific subroutines I have to do #1251 (comment) but this shouldn't be required for my scenarios because it makes more sense to jit the entire computation of F, so I'll just define:

fast_F = jax.jit(NonLinearF.compute) and then I have the correct functions with all subroutine properly defined, with well-organised code. If decorator is important you can use

I think it's a good middle point between jax functional philosophy (the cls from classmethod still doesn't have internal state if we don't define any static class attributes) and the convenience of inheritance for modular functionality: that's essentially a namespace with an object-oriented way of redefinng functions.

Still, If anyone has another way of organising code properly in Jax and without classes, I'm interested!

[patrick-kidger]

You may find Equinox interesting. It's framed as a "neural network library" but it's actually something a bit more general: a parameterised function library. In particular this gives a reasonable class-based/inheritance-supporting way of working with JAX.

As an example quite close to your own: Diffrax uses Equinox very extensively to define hierarchies of abstract parameterised functions.

For example, you can define:

import abc
import equinox as eqx
import jax.numpy as jnp

class AbstractFoo(eqx.Module):
    param: jnp.ndarray

    def bar(self, other):
        return self.param + other + self.baz()

    @abc.abstractmethod
    def baz(self):
        pass

class MyFoo(AbstractFoo):
    another_param: jnp.ndarray

    def baz(self):
        return self.another_param

Which you can use in the way you expect:

my_foo = MyFoo(jnp.array(1.0), jnp.array(2.0))
out = my_foo.bar(3)

and moreover can interop with jax.grad, jax.jit etc. without worry as well:

import jax

@jax.jit
@jax.grad
def call(foo, x):
    return foo.bar(x)

call(my_foo, 2.0)