jnp.linalg.norm error

[arvindmn01]

I am trying to finetune wav2vec2 model along with bart as language model using jax/flax. but I got this error

    main()
  File "run_flax_speech_recognition_seq2seq.py", line 1189, in main
    state, train_metric = p_train_step(state, batch)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/api.py", line 1779, in cache_miss
    execute = pxla.xla_pmap_impl_lazy(fun_, *tracers, **params)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 411, in xla_pmap_impl_lazy
    compiled_fun, fingerprint = parallel_callable(
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/linear_util.py", line 345, in memoized_fun
    ans = call(fun, *args)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 678, in parallel_callable
    pmap_computation = lower_parallel_callable(
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 825, in lower_parallel_callable
    jaxpr, consts, replicas, shards = stage_parallel_callable(pci, fun)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/interpreters/pxla.py", line 748, in stage_parallel_callable
    jaxpr, out_sharded_avals, consts = pe.trace_to_jaxpr_final(
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/interpreters/partial_eval.py", line 2233, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/interpreters/partial_eval.py", line 2177, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers_)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "run_flax_speech_recognition_seq2seq.py", line 1051, in train_step
    "encoder_grad_norm": jnp.linalg.norm(jax.tree_util.tree_leaves(layer_grad_norm["encoder"])),
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 166, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/pjit.py", line 250, in cache_miss
    outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/pjit.py", line 158, in _python_pjit_helper
    args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/api.py", line 306, in infer_params
    return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/pjit.py", line 505, in common_infer_params
    jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/pjit.py", line 971, in _pjit_jaxpr
    jaxpr, final_consts, out_type = _create_pjit_jaxpr(
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/linear_util.py", line 345, in memoized_fun
    ans = call(fun, *args)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/pjit.py", line 924, in _create_pjit_jaxpr
    jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/interpreters/partial_eval.py", line 2155, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/interpreters/partial_eval.py", line 2177, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers_)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/numpy/linalg.py", line 486, in norm
    check_arraylike("jnp.linalg.norm", x)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/numpy/util.py", line 328, in check_arraylike
    raise TypeError(msg.format(fun_name, type(arg), pos))
jax._src.traceback_util.UnfilteredStackTrace: TypeError: jnp.linalg.norm requires ndarray or scalar arguments, got <class 'list'> at position 0.

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 "run_flax_speech_recognition_seq2seq.py", line 1266, in <module>
    main()
  File "run_flax_speech_recognition_seq2seq.py", line 1189, in main
    state, train_metric = p_train_step(state, batch)
  File "run_flax_speech_recognition_seq2seq.py", line 1051, in train_step
    "encoder_grad_norm": jnp.linalg.norm(jax.tree_util.tree_leaves(layer_grad_norm["encoder"])),
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/numpy/linalg.py", line 486, in norm
    check_arraylike("jnp.linalg.norm", x)
  File "/home/arvind/wav2vec2_env/lib/python3.8/site-packages/jax/_src/numpy/util.py", line 328, in check_arraylike
    raise TypeError(msg.format(fun_name, type(arg), pos))
TypeError: jnp.linalg.norm requires ndarray or scalar arguments, got <class 'list'> at position 0.

jax version is jax==0.4.13
jaxlib version is jaxlib==0.4.13
flax version is flax==0.7.2
any help from anyone would be appreciated.

[jakevdp]

From the traceback, it seems the error is in line 1051 of File "run_flax_speech_recognition_seq2seq.py":

    "encoder_grad_norm": jnp.linalg.norm(jax.tree_util.tree_leaves(layer_grad_norm["encoder"])),

The output of tree_leaves is a list, and the input of jnp.linalg.norm must be an array. You could convert this list to an array to fix the error:

    "encoder_grad_norm": jnp.linalg.norm(jnp.asarray(jax.tree_util.tree_leaves(layer_grad_norm["encoder"]))),

[arvindmn01]

Hi @jakevdp thanks for reply.
I tried this code and it worked for me.
"encoder_grad_norm": jnp.linalg.norm(jax.tree_util.tree_leaves(layer_grad_norm["encoder"])[0]),