Can I write a Neural Network with Trax in a functional way?

Question

I am trying to learn Trax. I have previous exprience with Tensorflow, and I prefer writing neural networks with functional api(https://www.tensorflow.org/guide/keras/functional). I was wondering if it is possible to do the same with Trax, the resources I checked are in Serial(Sequential in Tenforflow/keras).

example code:

input = (placeholder)
layer1 = (placeholder)(input)
output = (placeholder)(layer1)