UML for deep learning architecture

Hello everyone,

can anybody share his experience with tools / methods for a clean documentation of Models?
I’ve tried using class diagram for my NTM.

Also I had the idea of drawing an activity diagram with a swimlane for each component to describe the flow of Tensors trough the model like this …