Stack Semantics

We're going to define a trax function (FN) for addition.

def Addition():
    layer_name = "Addition" 

    def func(x, y):
        return x + y

    return layers.Fn(layer_name, func)

Test it out.

add = Addition()
print(type(add))

<class 'trax.layers.base.PureLayer'>

print("name :", add.name)
print("expected inputs :", add.n_in)
print("promised outputs :", add.n_out)

name : Addition
expected inputs : 2
promised outputs : 1

x = numpy.array([3])
y = numpy.array([4])

print(f"{x} + {y} = {add((x, y))}")

[3] + [4] = [7]

def Multiplication():
    layer_name = "Multiplication"

    def func(x, y):
        return x * y

    return layers.Fn(layer_name, func)

Test it out.

mul = Multiplication()

The properties.

print("name :", mul.name)
print("expected inputs :", mul.n_in)
print("promised outputs :", mul.n_out, "\n")

name : Multiplication
expected inputs : 2
promised outputs : 1 

Some Inputs.

x = numpy.array([7])
y = numpy.array([15])
print("x :", x)
print("y :", y)

x : [7]
y : [15]

The Output

z = mul((x, y))
print(f"{x} * {y} = {mul((x, y))}")

[7] * [15] = [105]

serial = layers.Serial(
    Addition(), Multiplication(), Addition()
)

inputs = (numpy.array([3]), numpy.array([4]), numpy.array([15]), numpy.array([3]))

serial.init(shapes.signature(inputs))

print(serial, "\n")
print("name :", serial.name)
print("sublayers :", serial.sublayers)
print("expected inputs :", serial.n_in)
print("promised outputs :", serial.n_out, "\n")

Serial_in4[
  Addition_in2
  Multiplication_in2
  Addition_in2
] 

name : Serial
sublayers : [Addition_in2, Multiplication_in2, Addition_in2]
expected inputs : 4
promised outputs : 1 

print(f"{inputs} -> {serial(inputs)}")

(array([3]), array([4]), array([15]), array([3])) -> [108]

The example with the two simple adition and multiplication functions that where coded together with the serial combinator show how stack semantics work in Trax.

Suppose we want to make the simple calculation \((3 + 4) \times 3 + 4\). We can use Select to perform the calculations in the following manner:

1. input 3 4
2. tl.Select([0, 1, 0, 1])
3. add
4. mul
5. add.

The tl.Select requires a list or tuple of 0-based indices to select elements relative to the top of the stack. For our example, the top of the stack is 3 (which is at index 0) then 4 (index 1) and we us Select to copy the top two elements of the stack and then push all four elements back onto the stack which after the command executes will now contain 3 4 3 4. The steps of the calculation for our example are shown in the table below. As in the previous table each column shows the contents of the stack and the outputs after the operations are carried out.

After processing all the inputs the stack contains 25 which is the result of the calculations.

serial = layers.Serial(
    layers.Select([0, 1, 0, 1]),
    Addition(),
    Multiplication(),
    Addition()
)

Now we'll create the input.

x = (numpy.array([3]), numpy.array([4]))
serial.init(shapes.signature(x))

print(serial, "\n")
print("name :", serial.name)
print("sublayers :", serial.sublayers)
print("expected inputs :", serial.n_in)
print("promised outputs :", serial.n_out, "\n")

Serial_in2[
  Select[0,1,0,1]_in2_out4
  Addition_in2
  Multiplication_in2
  Addition_in2
] 

name : Serial
sublayers : [Select[0,1,0,1]_in2_out4, Addition_in2, Multiplication_in2, Addition_in2]
expected inputs : 2
promised outputs : 1 

print(f"{x} -> {serial(x)}")

(array([3]), array([4])) -> [25]

Note that since you are passing in indices to Select, you aren't really using it like a stack, even if behind the scenes it's using push and pop.

serial = layers.Serial(
    layers.Select([2, 1, 1, 2]),
    Addition(),
    Multiplication(),
    Addition()
)

x = (numpy.array([3]), numpy.array([4]), numpy.array([5]))
serial.init(shapes.signature(x))

print(f"{x} -> {serial(x)}")

(array([3]), array([4]), array([5])) -> [41]

print((5 + 4) * 4 + 5)

41

Residual networks (that link is to a research paper, this is wikipedia)are frequently used to make deep models easier to train. Trax already has a built in layer for this. The Residual layer computes the element-wise sum of the stack-top input with the output of the layer series. Let's first see how it is used in the code below:

serial = layers.Serial(
    layers.Select([0, 1, 0, 1]),
    layers.Residual(Addition())
)

print(serial, "\n")
print("name :", serial.name)
print("expected inputs :", serial.n_in)
print("promised outputs :", serial.n_out)

Serial_in2_out3[
  Select[0,1,0,1]_in2_out4
  Serial_in2[
    Branch_in2_out2[
      None
      Addition_in2
    ]
    Add_in2
  ]
] 

name : Serial
expected inputs : 2
promised outputs : 3

Here, we use the Serial combinator to define our model. The inputs first goes through a Select layer, followed by a Residual layer which passes the Fn: Addition() layer as an argument. What this means is the Residual layer will take the stack top input at that point and add it to the output of the Fn: Addition() layer. You can picture it like the diagram the below, where x1 and x2 are the inputs to the model:

Now, let's try running our model with some sample inputs and see the result:

x1 = numpy.array([3])
x2 = numpy.array([4])

print(f"{x1} + {x2} -> {serial((x1, x2))}")

[3] + [4] -> (array([10]), array([3]), array([4]))

As you can see, the Residual layer remembers the stack top input (i.e. 3) and adds it to the result of the Fn: Addition() layer (i.e. 3 + 4 = 7). The output of Residual(Addition() is then 3 + 7 = 10 and is pushed onto the stack.

On a different note, you'll notice that the Select layer has 4 outputs but the Fn: Addition() layer only pops 2 inputs from the stack. This means the duplicate inputs (i.e. the 2 rightmost arrows of the Select outputs in the figure above) remain in the stack. This is why you still see it in the output of our simple serial network (i.e. array([3]), array([4])). This is useful if you want to use these duplicate inputs in another layer further down the network.

To strengthen your understanding, you can modify the network above and examine the outputs you get. For example, you can pass the Fn: Multiplication() layer instead in the Residual block:

serial = layers.Serial(
    layers.Select([0, 1, 0, 1]), 
    layers.Residual(Multiplication())
)

print(serial, "\n")
print("name :", serial.name)
print("expected inputs :", serial.n_in)
print("promised outputs :", serial.n_out)

Serial_in2_out3[
  Select[0,1,0,1]_in2_out4
  Serial_in2[
    Branch_in2_out2[
      None
      Multiplication_in2
    ]
    Add_in2
  ]
] 

name : Serial
expected inputs : 2
promised outputs : 3

This means you'll have a different output that will be added to the stack top input saved by the Residual block. The diagram becomes like this:

And you'll get 3 + (3 * 4) = 15 as output of the Residual block:

x1 = numpy.array([3])
x2 = numpy.array([4])

y = serial((x1, x2))
print(f"{x1} * {x2} -> {serial((x1, x2))}")

[3] * [4] -> (array([15]), array([3]), array([4]))