How to use trax.layers.convolution.Conv1d for a simple 1D convolution operation? #question #help-wanted

[esmagden]

I'm having trouble passing arguments to the convolution layer for just a single evaluation of the layer. So instead of building a network, I'd just like to pass in the input and the kernel, and get the output from that, with plans to backpropagate through it afterwards. Clearly, there's something wrong about how the arguments are specified in my simple example. The following is running on Google CoLab.

import os
!pip install trax
import trax
from trax.fastmath import numpy as fastnp
trax.fastmath.use_backend('jax')
from trax import layers as tl

batch_size = 1
x = fastnp.ones([batch_size, 10, 1])*2
mykernel = tf.ones([4, 1, 1])

mylayer = tl.convolution.Conv1d(1, 4, padding='SAME', kernel_initializer='Ones', bias_initializer=None)
y = mylayer(x, weights = mykernel)

Something similar to this runs fine on TensorFlow, but I'm having trouble figuring out which inputs need to be specified to mylayer, and exactly how. It either says too many values to unpack, or not enough values. It's also not clear how to figure this out from the documentation, but any pointers on how to interpret that would also be helpful.

Exception passing through layer Conv (in pure_fn):
  layer created in file [...]/<ipython-input-71-2ddad8a38628>, line 2
  layer input shapes: ShapeDtype{shape:(1, 10, 1), dtype:float32}

  File [...]/trax/layers/convolution.py, line 56, in forward
    w, b = self.weights

ValueError: too many values to unpack (expected 2)

Edit: Added the labels to the title, since it won't allow for non-contributors to add them. Will remove when correctly labeled by the contributors.

[zack466]

Before using a layer or model, you need to initialize the weights. After defining mylayer, you would call mylayer.init(trax.shapes.signature(x)). Then, you can do mylayer(x) to get the output.

Also, for the kernel and bias initializers, trax is expecting functions, not strings/None, so that would also give you an error.

It would probably be helpful to read the layers intro and the quick intro, if you haven't already.