In the article, we propose an implementation of the matrix representation of Lie transform using TensorFlow as a computational engine. TensorFlow allows easy description of deep neural networks and provides automatic code execution on both single CPU/GPU and cluster architectures. In this research, we demonstrate the connection of the matrix Lie transform with polynomial neural networks. The architecture of the neural network is described and realized in code. In terms of beam dynamics, the proposed technique provides a tool for both simulation and analysis of experimental results using modern machine learning techniques. As a simulation technique one operates with a nonlinear map up to the necessary order of nonlinearity. On the other hand, one can utilize TensorFlow engine to run map optimization and system identification problems.