Improve comments

Author: Clifford-appliedAI-GmbH

prototree/train.py

@@ -84,8 +84,8 @@ def update_leaf_distributions(
     #      y_true_range = torch.arange(0, batch_size)
     #      y_true_indices = torch.stack((y_true_range, y_true))
     #      y_true_one_hot = torch.sparse_coo_tensor(y_true_indices,
-    #          torch.ones_like(y_true, dtype=torch.bool), logits.shape)  # Or other more suitable sparse format,
-    #  or even better,
+    #          torch.ones_like(y_true, dtype=torch.bool), logits.shape)  # Might be better to use CSR or CSC
+    #  or better still,
     #      y_true_one_hot = F.sparse_one_hot(y_true, num_classes=num_classes, dtype=torch.bool),
     #  but PyTorch doesn't yet have sufficient sparse mask support for the logic in update_leaf to work.
     y_true_one_hot = F.one_hot(y_true, num_classes=num_classes).to(dtype=torch.bool)
@@ -113,6 +113,12 @@ def update_leaf(
     log_p_arrival = node_to_prob[leaf].log_p_arrival.unsqueeze(1)
     # shape (num_classes). Not the same as logits, which has (batch_size, num_classes)
     leaf_logits = leaf.logits()
+
+    # TODO If PyTorch had more support for sparse masks we might be able to do something like
+    #      masked_logits = logits.sparse_mask(y_true_one_hot),
+    #  and perhaps if necessary combine it with
+    #      masked_log_p_arrival = y_true_one_hot * log_p_arrival  # sparse_mask can't broadcast
+    #      masked_leaf_logits = y_true_one_hot * leaf_logits  # sparse_mask can't broadcast.
     masked_logits = masked_tensor(logits, y_true_one_hot)
 
     masked_log_combined = log_p_arrival + (leaf_logits - masked_logits)