PyTorch implementation of soft decision tree node. This recursive version is more readable however the tensorized version is faster.

SoftNode.py

import torch

class SoftNode(torch.nn.Module):
    def __init__(self, in_features, out_features, depth, projection="constant"):
        super(SoftNode, self).__init__()
        self.projection = projection
        if depth > 0:
            self.left = SoftNode(in_features, out_features, depth-1, projection=projection)
            self.right = SoftNode(in_features, out_features, depth-1, projection=projection)
            self.gating = torch.nn.Linear(in_features, 1)
            self.expert = None
        else:
            self.left = None
            self.right = None
            self.gating = None
            if projection == "constant":
                self.expert = torch.nn.Parameter(torch.randn(out_features))
            elif projection == "linear":
                self.expert = torch.nn.Linear(in_features, out_features)
    
    def forward(self, x):
        if self.expert is None:
            g = torch.sigmoid(self.gating(x))
            out = g * self.left(x) + (1-g) * self.right(x)
        else:
            if self.projection == "constant":
                out = self.expert
            else:
                out = self.expert(x)
        return out