Issue with DiscreteJaxOperator

[XiaoweiOu]

After updating to the latest version, I encountered an error when calling get_conn_flattened function of the DiscreteJaxOperator class. Line 171 tries to convert xp to a Numpy array, where I got error message:

jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method array() was called on traced array with shape int8[1,1,8]
This BatchTracer with object id 22743392749904 was created on line:
/home/xo24/SSH_model/venv-jax/lib64/python3.11/site-packages/netket/utils/static_range.py:192:16 (StaticRange.numbers_to_states)
See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerArrayConversionError

I don't understand why we convert traced variables to Numpy arrays in DiscreteJaxOperator class. This seems to contradict with the jax function transformation. Could anyone explain this update? Thanks for any help!

[PhilipVinc]

ah, indeed. it's a bug. Can easily be fixed.

[PhilipVinc]

Do you have a code that errors? A script? In theory it should not crash in normal usage?

[XiaoweiOu]

This is the bosonic operator that I use:

from netket.operator.boson import number as na
# Phonon number hamiltonian
for i in range(boson_hilbert.size):
    hamil_phonon_number += na(boson_hilbert, i)

Then, I convert the operator to the JAX version.

boson_op_jax = hamil_phonon_number.to_jax_operator()

After that, in another function wrapped by the decorator equinox.filter_jit, this line causes the error:

phonon_conn_states, phonon_conn_mels = boson_op_jax.get_conn(boson_state)

Here I only process one configuration, so I use get_conn function. Later I use jax.vmap to process batched input.

Please let me know if this code snippet is not enough. Thank you!

[PhilipVinc]

Can you provide a full MWE, that is, a reasonably short script that I can copy-paste and which showcases the error you are talking about?

[PhilipVinc]

In any case, I think the issue is that get_conn cannot be vmapped because it has dynamic shape sizes.
If you need to do it on batched inputs (or need to vmap) you should use get_conn_padded

[XiaoweiOu]

Oh yes! It works after I change to get_conn_padded. Now I noticed 'get_conn_padded method can be executed inside of a Jax function transformation.' in the documentation.

Thanks so much for your help!

[PhilipVinc]

Thanks for reporting it. I added a more informative error now when you attempt to do what you were doing. It will clearly suggest to use get_conn_padded.

[XiaoweiOu]

Thank you for your reply.

For now, can I define a customized function to circumvent this problem? Or should I wait for the next release?