Creating a GRU model using Trax: Ungraded Lecture Notebook

For this lecture notebook you will be using Trax's layers. These are the building blocks for creating neural networks with Trax.

import trax
from trax import layers as tl

INFO:tensorflow:tokens_length=568 inputs_length=512 targets_length=114 noise_density=0.15 mean_noise_span_length=3.0

Trax allows to define neural network architectures by stacking layers (similarly to other libraries such as Keras). For this the Serial() is often used as it is a combinator that allows to stack layers serially using function composition.

Next you can see a simple vanilla NN architecture containing 1 hidden(dense) layer with 128 cells and output (dense) layer with 10 cells on which we apply the final layer of logsoftmax.

mlp = tl.Serial(
  tl.Dense(128),
  tl.Relu(),
  tl.Dense(10),
  tl.LogSoftmax()
)