Need help using jax.grad -- getting type error in loss function

[timnewsham]

Hi, I need some help using jax.grad. I'm getting a type error in my loss function, defined as:

def square_error(expected, actual):
    print(repr(expected), repr(actual))
    error = expected - actual
    return error.dot(error)

def batch_loss(layers, training_data):
    return sum(
        square_error(expected_output, network(input, layers))
        for (input, expected_output) in training_data
    )

def gradient(n, loss, layers, training_data):
    loss0 = batch_loss(layers, training_data)
    print(f"{n} Loss: {loss0}\r", end='')
    return jax.grad(loss)(layers, training_data)

when computing error = expected - actual the expected field is a jax array and the actual field seems to be some wrapped type that the gradient code is using to track data.

The entire project (not large) is in the attached file, and invoked with python test_train.py.
jax_help.tgz

The error I received is:

Traceback (most recent call last):
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/./test_train.py", line 21, in <module>
    test()
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/./test_train.py", line 12, in test
    net = train((2,2), square_error, training_data, iters=1000)
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/train.py", line 45, in train
    gradient_descent(n, loss, layers, training_data, factor)
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/train.py", line 24, in gradient_descent
    grads = gradient(n, loss, layers, training_data)
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/train.py", line 21, in gradient
    return jax.grad(loss)(layers, training_data)
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 163, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/api.py", line 1176, in grad_f
    _, g = value_and_grad_f(*args, **kwargs)
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 163, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/api.py", line 1252, in value_and_grad_f
    ans, vjp_py = _vjp(f_partial, *dyn_args, reduce_axes=reduce_axes)
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/api.py", line 2794, in _vjp
    out_primal, out_vjp = ad.vjp(
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 137, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/interpreters/ad.py", line 126, in linearize
    jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals)
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 770, in trace_to_jaxpr_nounits
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/linear_util.py", line 165, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/train.py", line 9, in square_error
    error = expected - actual
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 5073, in deferring_binary_op
    raise TypeError(f"unsupported operand type(s) for {opchar}: "
jax._src.traceback_util.UnfilteredStackTrace: TypeError: unsupported operand type(s) for -: 'list' and 'Array'

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/./test_train.py", line 21, in <module>
    test()
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/./test_train.py", line 12, in test
    net = train((2,2), square_error, training_data, iters=1000)
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/train.py", line 45, in train
    gradient_descent(n, loss, layers, training_data, factor)
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/train.py", line 24, in gradient_descent
    grads = gradient(n, loss, layers, training_data)
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/train.py", line 21, in gradient
    return jax.grad(loss)(layers, training_data)
  File "/Users/timnewsham/src/pyNeuralNet/with_jax/x/train.py", line 9, in square_error
    error = expected - actual
  File "/Users/timnewsham/.pyenv/versions/3.10.8/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 5073, in deferring_binary_op
    raise TypeError(f"unsupported operand type(s) for {opchar}: "
TypeError: unsupported operand type(s) for -: 'list' and 'Array'

[jakevdp]

Can you include the traceback of the error you're seeing?

Just a guess without running your code; perhaps you need jnp.sum rather than sum in batch_loss

[timnewsham]

I just realized I forgot to included. I've updated the original question to include it.

[jakevdp]

It looks like your training data contains lists rather than arrays. Try converting expected_output to an array.

[timnewsham]

My training data in the stack trace is:

(Pdb) print(repr(training_data))
Array([[[0., 0.],
        [1., 0.]],

       [[0., 1.],
        [1., 0.]],

       [[1., 0.],
        [1., 0.]],

       [[1., 1.],
        [0., 1.]]], dtype=float32)

all values of expected should be arrays. When the square_error is called "normally" with the training data it always has the right types. I only see it having this weird type when its called from jax.grad.

I'm using

Python 3.10.8 (main, Dec  5 2022, 09:36:08) [Clang 14.0.0 (clang-1400.0.29.202)] on darwin
>>> jax.__version__
'0.4.4'

$ uname -a
Darwin TimBookPro.lan 22.3.0 Darwin Kernel Version 22.3.0: Mon Jan 30 20:42:11 PST 2023; root:xnu-8792.81.3~2/RELEASE_X86_64 x86_64

[jakevdp]

According to the traceback, expected is a list rather than an array. You should make sure it's an array.

[jakevdp]

Or it may be that network returns a list rather than an array

[timnewsham]

network returns values like Array([9.999938e-01, 6.144174e-06], dtype=float32).
I call the same loss function both without and with jax.grad. The call to the loss function without jax.grad works just fine:

def gradient(n, loss, layers, training_data):
    loss0 = batch_loss(layers, training_data)     # <--- works
    print(f"{n} Loss: {loss0}\r", end='')
    return jax.grad(loss)(layers, training_data)  # <--- fails

Adding some prints to the square_error function reports the proper types when calling doesnt go through jax.grad:

def square_error(expected, actual):
    print(repr(expected), repr(actual))
    error = expected - actual
    return error.dot(error)

giving

Array([1., 0.], dtype=float32) Array([0.37566873, 0.6243313 ], dtype=float32)
Array([1., 0.], dtype=float32) Array([0.43291855, 0.56708145], dtype=float32)
Array([1., 0.], dtype=float32) Array([0.54194355, 0.45805645], dtype=float32)
Array([0., 1.], dtype=float32) Array([0.6338155 , 0.36618447], dtype=float32)

When called through jax.grad the print statement gives:

[(Traced<ConcreteArray([[0.9653214  0.22515893]
 [0.63302994 0.29638183]], dtype=float32)>with<JVPTrace(level=2/0)> with
  primal = Array([[0.9653214 , 0.22515893],
       [0.63302994, 0.29638183]], dtype=float32)
  tangent = Traced<ShapedArray(float32[2,2])>with<JaxprTrace(level=1/0)> with
    pval = (ShapedArray(float32[2,2]), None)
    recipe = LambdaBinding(), Traced<ConcreteArray([0.21629536 0.8041241 ], dtype=float32)>with<JVPTrace(level=2/0)> with
  primal = Array([0.21629536, 0.8041241 ], dtype=float32)
  tangent = Traced<ShapedArray(float32[2])>with<JaxprTrace(level=1/0)> with
    pval = (ShapedArray(float32[2]), None)
    recipe = LambdaBinding())] Array([[[0., 0.],
        [1., 0.]],

       [[0., 1.],
        [1., 0.]],

       [[1., 0.],
        [1., 0.]],

       [[1., 1.],
        [0., 1.]]], dtype=float32)

[timnewsham]

ahh! I see a problem.. I was calling batch_loss not loss, and I wasnt passing in the right loss function!!!
I think I have this resolved.