1
class Einsum(nn.Module):
2
    """Einsum with LoRA support. Can be used as a drop-in replacement for the Gemma Einsum."""
3

4
    # Shape of the weight.
5
    shape: tuple[int, ...]
6
    # Initialization function for the weight.
7
    init_fn: nn.initializers.Initializer = nn.initializers.zeros
8
    # If not None, apply LoRA to the weight.
9
    lora_config: LoRAConfig | None = None
10

11
    def setup(self):
12
        self.w = self.param("w", self.init_fn, self.shape)
13

14
        if config := self.lora_config:
15
            # Setup LoRA parameters.
16
            shape_a, shape_b = list(self.shape), list(self.shape)
17
            shape_a[config.axes[1]] = config.rank
18
            shape_b[config.axes[0]] = config.rank
19
            self.w_a = self.param("lora_a", config.init_fn, shape_a)
20
            self.w_b = self.param("lora_b", config.init_fn, shape_b)
21

22
    @nn.compact
23
    def __call__(self, eqn: str, x):
24
        dtype = x.dtype  # original dtype, could be half-precision
25
        result = jnp.einsum(eqn, x, self.w.astype(dtype))
26

27
        if config := self.lora_config:
28
            eqn_a, eqn_b = self._make_lora_eqns(eqn)
29
            lora = jnp.einsum(eqn_a, x, self.w_a.astype(dtype))
30
            lora = jnp.einsum(eqn_b, lora, self.w_b.astype(dtype))
31
            result = result + lora * config.scaling_value
32

33
        return result