implementing a neural network manually in Jax.

[HosseinKhodavirdi]

Hi, I am new to Jax and I am trying to solve a deep learning problem by implementing its network manually. For instance, this part, Computes the forward pass for each example individually: (np is jax numpy here)

def NN(params, r, ang):
    # print(x.shape,y.shape,k.shape)
    activations = np.array([r, ang])
    # Loop over the ReLU hidden layers
    for w, b in params[:-1]:
        activations = relu3_layer([w, b], activations)
    final_w, final_b = params[-1]
    final = np.sum(np.dot(final_w, activations)) + final_b
    return (final[0])

I have multiple questions which I appreciate any help:

1. During the training, based on a need I tried to print the out put of NN. Based on my problem I do expect to have an array with the shape of r or ang in the output. Despite, this is what I get at a few first epochs:

Traced<ShapedArray(float32[])>with<JVPTrace(level=9/0)> with
  primal = Traced<ShapedArray(float32[])>with<JVPTrace(level=7/0)> with
    primal = Traced<ShapedArray(float32[])>with<BatchTrace(level=5/0)> with
      val = Traced<ShapedArray(float32[162])>with<DynamicJaxprTrace(level=4/0)>
      batch_dim = 0
    tangent = Traced<ShapedArray(float32[])>with<BatchTrace(level=6/0)> with
      val = Traced<ShapedArray(float32[1])>with<BatchTrace(level=5/0)> with
        val = Traced<ShapedArray(float32[162,1])>with<DynamicJaxprTrace(level=4/0)>
        batch_dim = 0
      batch_dim = 0
  tangent = Traced<ShapedArray(float32[])>with<BatchTrace(level=8/0)> with
    val = Traced<ShapedArray(float32[1])>with<JVPTrace(level=7/0)> with
      primal = Traced<ShapedArray(float32[1])>with<BatchTrace(level=5/0)> with
        val = Traced<ShapedArray(float32[162,1])>with<DynamicJaxprTrace(level=4/0)>
        batch_dim = 0
      tangent = Traced<ShapedArray(float32[1])>with<BatchTrace(level=6/0)> with
        val = Traced<ShapedArray(float32[1,1])>with<BatchTrace(level=5/0)> with
          val = Traced<ShapedArray(float32[162,1,1])>with<DynamicJaxprTrace(level=4/0)>
          batch_dim = 0
        batch_dim = 0
    batch_dim = 0

• I really want to know why I have a 9 level traced array? What does this level mean and why it is 9? Also why in the first line it is showing ShapedArray(float32[]), while based on my expectation, it should be Traced<ShapedArray(float32[162]).
• Now let's say I have a function in the middle of my code which uses the output of NN and tries to implement an element-by-element multiplication with another array at each epoch. I assume based on the above script the output of NN is either empty or very complicated which does not let me to do an element-by-element multiplication with another array:

temp=NN(params,r,ang)
temp2=np.multiply(cosnp, temp)

• The thing that surprises me more is that when the code keeps running, the output of NN first changes to:

 Traced<ShapedArray(float32[])>with<BatchTrace(level=5/0)> with
  val = Traced<ShapedArray(float32[162])>with<DynamicJaxprTrace(level=4/0)>
  batch_dim = 0

and then to:

Traced<ShapedArray(float32[])>with<BatchTrace(level=1/0)> with
  val = Array([-1.7656987e-09,  0.0000000e+00, -1.7444961e-09, -1.7159607e-09,
       -1.6791830e-09, -1.6331967e-09, -1.5770922e-09, -1.5098749e-09,
       -1.4304362e-09, -1.3379235e-09, -1.2311432e-09, -1.1089867e-09,
       -9.7014663e-10, -8.1351459e-10, -6.3764105e-10, -4.4119020e-10,
       -2.2296831e-10,  1.8474111e-11,  2.8447289e-10,  5.7625016e-10,
        8.9525543e-10,  1.2428529e-09,  1.6201795e-09,  2.0287132e-09,
        2.4698181e-09,  2.9447449e-09,  3.4549714e-09,  4.0017767e-09,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  0.0000000e+00,
        0.0000000e+00,  0.0000000e+00,  0.0000000e+00], dtype=float32)
  batch_dim = 0

Can someone explain what do all these things mean?

My last question is about using the grad operator in JAX. I have a function which its output may be complex in some epochs. I tried grad to take the derivative of my function but it showed an error that it does not work for complex values. Setting holomorphic=True also does not work since the function do not have always a complex type. Does anyone have any suggesiton?

Thanks much in advance,

[jakevdp]

Hi - it sounds like you're curious about what JAX tracers are, and how you can think of the execution model for JAX programs. I'd start with these docs:

• How To Think In JAX
• FAQ: Different Kinds of JAX Values

To more specifically answer your questions, I wouldn't worry about the level number in "Level 9 tracer". When you see a tracer, it just means that your code is being transformed by a JAX transformation like jit or grad. If you're interested in printing intermediate traced values, check out jax.debug.print.

JAX functions should work correctly with traced arrays, which is why you see that further computations work as expected despite the fact that the inputs are tracers.

Hope that helps!

[HosseinKhodavirdi]

Thank Jake. What if I need to apply a jax convolution on a traced variable. The point of using jnp.convolve is that it needs 1-D array as input while the shape of traced variable is ( ). How can such a thing be done?

Traced<ShapedArray(float32[])>with<JVPTrace(level=9/0)> with

[jakevdp]

That can happen, for example, if you are calling jax.vmap on your function: then a 1-dimensional input would result in a 0-dimensional tracer.

[HosseinKhodavirdi]

Sorry, I think I did not explain well. I have a traced array like (NN):

print(NN) =====>
Traced<ShapedArray(float32[])>with<JVPTrace(level=9/0)> with...

and I want to calculate convolution integral of that traced array with another array using jnp.convolve. The jnp convolve needs 1D arrays as input while here the traced array I have (NN) shows a shape of ( ). How should I handle such situations. Could not find the answer to it in docs.

[jakevdp]

You cannot call jnp.convolve on zero-dimensional arrays. You should change how you call the function and instead pass a one-dimensional array into the function.

If you are already passing a one-dimensional array into the function, then it's likely that your function is wrapped with vmap or another transformation that causes it to effectively be evaluated element-wise.

[jakevdp]

It's hard to guess what's happening here because the code you shared doesn't even reference jnp.convolve – you might try editing your question to add a minimal reproducible example if you're looking for more specific help.