How to print the value in jax

[jimmylihui]

The print doesn't print the information of error. Does anyone know how to print in Jax?

[jakevdp]

Hi - thanks for the question!

In the future, please copy and paste the code as text rather than as a screenshot – you can put it in a block between triple back-tick-marks (```) to format it. It's also helpful if you construct a minimal reproducible example of the issue you're having, so we have to make fewer guesses about what might be causing the problem.

To answer your question, the call to jax.debug.print looks correct to me, and should print the value contained in error[0] to stdout. One guess: if you're running your code in a Jupyter notebook, the stdout stream may be displayed in the terminal from which the notebook was launched rather than in the output cell in the notebook.

Can you say more about how you're running your code, what you're actually seeing, and what you expect to see?

[jimmylihui]

@eqx.filter_jit
def Linear(x,y,parameter):
    alpha=0.1
    x=jnp.atleast_2d(x)
    #x[batch, input_dimension], y[batch, output dimension], parameter[input, output]
    forward=einsum('bi,io->bo',x,parameter)
    forward=sigmoid(forward)
    error=forward-y

    D=error*sigmoid_deriv(forward)
    # delta=D.dot(parameter.T)
    a=x.T.dot(D)
    jax.debug.print("parameter before:{parameter}", parameter=parameter[0])
    parameter+= -alpha*x.T.dot(D)
    jax.debug.print("parameter after:{parameter}", parameter=parameter[0])
    return parameter,error

The result is shown below

parameter before:[-1.7753183]
parameter after:[-1.7547464]
parameter before:[-1.7546065]
parameter after:[-1.7546065]
parameter before:[-1.7633415]
parameter after:[-1.7633415]
parameter before:[-1.7656941]
parameter after:[-1.7656941]
parameter before:[-1.7660894]
parameter after:[-1.7660894]
parameter before:[-1.8295524]
parameter after:[-1.8295524]
parameter before:[-1.8497633]
parameter after:[-1.8497633]
parameter before:[-1.8251902]
parameter after:[-1.8251902]
parameter before:[-1.8305788]
parameter after:[-1.8305788]
parameter before:[-1.8280289]
parameter after:[-1.8280289]
parameter before:[-1.7860788]
parameter after:[-1.7860788]
parameter before:[-1.798911]
parameter after:[-1.798911]
parameter before:[-1.7968287]
parameter after:[-1.7968287]
parameter before:[-1.7993668]
parameter after:[-1.7993668]
parameter before:[-1.7969431]
parameter after:[-1.7969431]
parameter before:[-1.7692602]
parameter after:[-1.7692602]

It seems backward propagation is not working.

[jakevdp]

If you want help with model correctness, we’ll need to see all the code. I’d suggest putting together a minimal reproducible example that demonstrates the behavior you’re seeing (see the link I shared in my initial response).

[jimmylihui]

import argparse
import time
import functools
import jax
import jax.numpy as jnp
import equinox as eqx

@eqx.filter_jit
def Linear(x,y,parameter):
    alpha=0.1
    x=jnp.atleast_2d(x)
    #x[batch, input_dimension], y[batch, output dimension], parameter[input, output]
    forward=einsum('bi,io->bo',x,parameter)
    forward=sigmoid(forward)
    error=forward-y

    D=error*sigmoid_deriv(forward)
    # delta=D.dot(parameter.T)
    a=x.T.dot(D)
    jax.debug.print("parameter before:{parameter}", parameter=parameter[0])
    parameter+= -alpha*x.T.dot(D)
    jax.debug.print("parameter after:{parameter}", parameter=parameter[0])
    return parameter,error

@functools.partial(jax.jit, static_argnames=( "model"))
def loop(parameter,input,output,model):
    error_list=[]
    for i in range(input.shape[0]):
        x=input[i:i+1,:]
        y=output[i:i+1,:]
        parameter,error=model(x,y,parameter)
        error_list.append(error)
        # jax.debug.print("{error}", error=error[0])
    return parameter,error_list


key = jax.random.PRNGKey(seed)
key, *subkey = jax.random.split(key, 4)
x = jax.random.normal(subkey[0], (10, 2), dtype=dtype)
y=jax.random.normal(subkey[1], (10, 1), dtype=dtype)
parameter=jax.random.normal(subkey[2], (2,1), dtype=dtype)
model=linear
parameter,error=loop(parameter,inputs,outputs,model)

please try this.

[jakevdp]

seed, dtype, and linear are undefined in your code.

[jimmylihui]

I have found the issue. Thanks for your help.