The Tensor type in Pytorch is a bit like a indexed monad with a value type, with the index being the tensor dimension. When a user compose multiple tensors through operators like multiplication and addition, it is similar to using an applicative functor. When a user performs loop with a tensor, it is similar to the monadic bind operation with the element values as the callback.
The main difference between the PyTorch's Tensor type and monads is that they also provide white box methods to "peek" and read the values inside through eager mode.
However, when PyTorch is used in lazy evaluation mode for advanced use cases, such as code generation based on the final tensor, then the whole program must be treated as a black box. With that, operations such as computation on the inner tensors must be done through a bind-like continuation-passing style interface.
To implement meta features such as code generation, PyToch is essentially treating the Tensor type like a free monad, so that the full program can be "interpreted" differently instead of being executed right away.
Tt might also be simpler to structure an ML program from a final encoding perspective. A lot of ML operations are actually just repeat of smaller primitive operations in a certain way.
For example, we can define type-indexed primitive operator of multiplication of two matrices of certain direction. We then also define higher-order repeat operators that turn the inner operator into an operator that takes in multi-dimensional matrices. There would also be operators like back propagation that reverse the input and output. The idea is that with final encoding, even the multi-dimensional matrices are still plain data types, and that we dont need to turn them into fancy monads.