NNT Tensor Operations

This is the node that turns ComfyUI into a whiteboard for calculus. NNT Tensor Operations takes one or two tensors, runs a math operation on them, and - if you pick a gradient op - shows you the derivative as a tensor you can inspect. It's the tool for anyone using this pack to see what gradients actually are, instead of just trusting that backprop happens somewhere inside the training loop.

It has a wider job list than that, though. Fourteen operations are on the menu: add, subtract, element-wise multiply, matrix multiply, transpose, inverse, add a scalar, multiply by a scalar, plus custom_function, custom_function_with_grad, gradient, jacobian, and a couple more. The boring ones are handy tensor plumbing; the gradient ones are where the pack earns its "educational toolkit" label.

How it works

The plumbing ops are straightforward torch calls on tensor_a and (optionally) tensor_b. The interesting part is the custom-function path: you type an expression into custom_expression (default tensor_a * 2) and the node evaluates it with tensor_a, tensor_b, and torch in scope. It handles a few special cases directly - tensor_a ** n becomes torch.pow, for example - and falls back to eval() for the rest, with builtins locked down.

The gradient machinery is where it shines:

• gradient - computes the gradient of the result with respect to the tensor you pick via grad_tensor (either tensor_a or tensor_b). It sums the result to a scalar and calls torch.autograd.grad, so you get one gradient per input element. Wire tensor_a, run gradient, and read the gradient_tensor output to see how each element would nudge the result.
• custom_function_with_grad - same idea, but the expression is yours, so you can differentiate tensor_a ** 3 - 2 * tensor_a without writing any math by hand.

Outputs are result_tensor, gradient_tensor (empty unless you asked for a gradient op), and info_message describing what ran.

What you'd use it for

Concrete use: learn what a derivative is. Set tensor_a to a scalar-ish tensor from the Random Tensor Generator, type tensor_a ** 2, run the gradient op, and read the output - it's 2·a, exactly as the algebra says. Then try tensor_a ** 3 and see 3·a². It's the "toy autograd" experiment, and it lands better visually than any textbook. The matrix-multiply and transpose ops also make it a reasonable mini linear-algebra checker while you're designing layer stacks.

A couple of honest caveats. The expression evaluator is a limited language - don't expect Python's full syntax, and if an expression fails it raises a ValueError you'll see in the console. Non-float tensors get cast to float for gradient work, so int inputs lose their dtype. And it needs requires_grad to be on for the tensor you differentiate against - the node will set it itself if you forgot, but that's the kind of invisible magic that surprises people later.

Install

Pack-level install:

cd ComfyUI/custom_nodes
git clone https://github.com/inventorado/ComfyUI_NNT.git
cd ComfyUI_NNT
pip install -r requirements.txt

Restart ComfyUI, or install via Manager by searching "ComfyUI Neural Network Toolkit NNT". Requirements are heavy (torch, numpy, matplotlib, sklearn, transformers, statsmodels, onnx, shap pinned at 0.41.0, etc.) so budget time for the first install. Worth it if "what even is a gradient" is the question you're here to answer.