Composable modules

[danjenson]

Hello,

I'm trying to create a composable module so that users of my library can do ergonomic operations like bias = DistanceBias(channel=0) + DistanceBias(channel=1) which will create a composite distance bias module. However, I can't seem to get parent / scoping sorted out and I always get a naming conflict. Here is an mvp of the code:

import operator
from dataclasses import dataclass
from typing import Callable, Optional, Union

import flax.linen as nn


@dataclass
class ExpressionNode:
    """Represents a computation tree for module compositions with correct precedence."""

    left: Union[nn.Module, float, "ExpressionNode"]
    op: Optional[Callable] = None
    right: Optional[Union[nn.Module, float, "ExpressionNode"]] = None

    def evaluate(self, x, variables):
        """Recursively evaluates the computation tree while supporting parameter updates."""
        if self.op is None:
            return (
                self.left.apply(variables, x)
                if isinstance(self.left, nn.Module)
                else self.left
            )

        left_val = (
            self.left.evaluate(x, variables)
            if isinstance(self.left, ExpressionNode)
            else self.left.apply(variables, x)
        )
        right_val = (
            self.right.evaluate(x, variables)
            if isinstance(self.right, ExpressionNode)
            else self.right.apply(variables, x)
            if isinstance(self.right, nn.Module)
            else self.right
        )

        return self.op(left_val, right_val)

    def __add__(self, other):
        return ExpressionNode(self, operator.add, other)

    def __sub__(self, other):
        return ExpressionNode(self, operator.sub, other)

    def __mul__(self, other):
        return ExpressionNode(self, operator.mul, other)

    def __truediv__(self, other):
        return ExpressionNode(self, operator.truediv, other)

    def __matmul__(self, other):
        return ExpressionNode(self, operator.matmul, other)

    def __neg__(self):
        return ExpressionNode(self, operator.mul, -1)

    def __rmul__(self, other):
        return ExpressionNode(self, operator.mul, other)


class ComposableModule(nn.Module):
    """A Flax module that supports arithmetic operations."""

    expr: Optional[ExpressionNode] = None

    @nn.compact
    def __call__(self, x):
        return self.expr.evaluate(x, self.variables)

    def _binary_op(self, other, op):
        left = self.expr if self.expr is not None else ExpressionNode(self)
        right = other.expr if other.expr is not None else ExpressionNode(other)
        return ComposableModule(expr=ExpressionNode(left, op, right))

    def __add__(self, other):
        return self._binary_op(other, operator.add)

    def __sub__(self, other):
        return self._binary_op(other, operator.sub)

    def __mul__(self, other):
        return self._binary_op(other, operator.mul)

    def __truediv__(self, other):
        return self._binary_op(other, operator.truediv)

    def __matmul__(self, other):
        return self._binary_op(other, operator.matmul)

    def __neg__(self):
        return self._binary_op(-1, operator.mul)

    def __rmul__(self, other):
        return self._binary_op(other, operator.mul)


class DistanceBias(ComposableModule):
    num_heads: int = 4
    channel: int = 0

    @nn.compact
    def __call__(
        self,
        d: jax.Array,  # [B, Q, K, D] or [E, D]
        mask: Optional[jax.Array] = None,  # None or [B, Q, K] or [E]
    ):
        a = self.param("a", init.constant(-1), (self.num_heads,))
        if mask is None:
            mask = jnp.array([True])
        return distance_bias(d[..., self.channel], mask, a)

When I run, DistanceBias(channel=0) + DistanceBias(channel=1), I get the following error:

File ~/.pyenv/versions/3.12.7/envs/jax/lib/python3.12/site-packages/flax/linen/module.py:1360, in Module.__post_init__(self)
  1355 reuse_scopes = (
  1356   self.parent._state.in_setup
  1357   and self.parent._state.setup_called == SetupState.TRANSFORMED
  1358 )
  1359 # Perform name-collision check.
-> 1360 if self.parent._name_taken(self.name, reuse_scopes=reuse_scopes):
  1361   parent_class = self.parent.__class__.__name__
  1362   raise errors.NameInUseError('submodule', self.name, parent_class)

File ~/.pyenv/versions/3.12.7/envs/jax/lib/python3.12/site-packages/flax/linen/module.py:1523, in Module._name_taken(self, name, reuse_scopes, collection)
  1517 def _name_taken(
  1518   self,
  1519   name: str,
  1520   reuse_scopes: bool = False,
  1521   collection: str | None = None,
  1522 ) -> bool:
-> 1523   assert self.scope is not None
  1524   if reuse_scopes:
  1525     return False

AssertionError:

Is there any reasonable way to do this in JAX?

[mattjj]

This might be a flax-specific question. If so, can you ask on the flax repo?

[danjenson]

Absolutely, apologies for the bother.

[mattjj]

No worries at all!