Integration of PQuantML with hls4ml

[enlupi]

Description

This PR introduces the possibility of directly parsing PQuantML layers into hls4ml, both using the PyTorch and Keras V3 frontend. It uses HGQ2 conventions and relies on the bit_exact optimization pass to ensure correct precision is enforced and bit exactness.

Type of change

• New feature (non-breaking change which adds functionality)

Tests

Tests are present both for PyTorch and Keras V3 frontend: test_pquant_pytoch.py and test_pquant_keras.py.

Checklist

• I have read the guidelines for contributing.
• I have commented my code, particularly in hard-to-understand areas.
• I have made corresponding changes to the documentation.
• My changes generate no new warnings.
• I have installed and run pre-commit on the files I edited or added.
• I have added tests that prove my fix is effective or that my feature works.

[JanFSchulte]

What's the plan for installing PQuant together with hls4ml? In principle, we should have pquant-ml defined as an optional dependency in the pyproject.toml. I see that the tool is available via pypi (https://libraries.io/pypi/pquant-ml), but just doing pip install pquant-ml in an empty conda environment with python 3.11 doesn't work for me because of version conflicts that prevent tensorflow from being installed. So I'm not sure the version available via pypi is up to date.

I expect the pytests to fail currently because of this. We are preparing a dedicated testing container with keras 3 that should be able to be used for this going forward, but that would then need the proper installation instructions for pquant to work.

Also, would be nice to have the documentation update included in this PR ;)

[enlupi]

What's the plan for installing PQuant together with hls4ml? In principle, we should have pquant-ml defined as an optional dependency in the pyproject.toml.

I have now included it.

doing pip install pquant-ml in an empty conda environment with python 3.11 doesn't work for me because of version conflicts that prevent tensorflow from being installed.

This is a known issue for the PQuantML team and they are trying to solve it. As far as I have understood, the problem is due to pip trying to install the wrong version of tensorflow, and the only solution at the moment is "manually" installing it with conda.

Also, would be nice to have the documentation update included in this PR ;)

Added :)

[JanFSchulte]

To make the tests work, pquant-ml also need to be added to the testing environments. Since it is not compatible with keras 2, I guess it should go into the optional dependencies for the keras 3 version of the environment https://github.com/fastmachinelearning/hls4ml/blob/main/pyproject.toml#L71 and then get added to the tests that use this environment: https://github.com/fastmachinelearning/hls4ml/blob/main/test/pytest/generate_ci_yaml.py#L31

We probably also need to run the torch version of the tests in this environment, then.

Locally I still have a lot of issues running the tests. I can install pquant through pip now, but I'm getting tons of errors coming from PQuant itself.

So I assume there is still something wrong with my environment.

[JanFSchulte]

I finally got the tests to work locally, so here is some feedback that should hopefully make it all work:

• The PQuant test using keras3 needs to be added to the keras3 test list here https://github.com/fastmachinelearning/hls4ml/blob/main/test/pytest/generate_ci_yaml.py#L31
• pquant-ml needs to be added to the dependencies installed for the test environement at https://github.com/fastmachinelearning/hls4ml/blob/main/pyproject.toml#L54

Even with that I still had some issues with pquant imports failing when running pytest, despite the same inputs working fine in a simple python shell in the same environment. But it worked when running it as python -m pytest. My hope is that it will work fine in the proper test environment, but lets see.

Then, running the actual tests, I get the following to types of errors:

• In the keras3 tests, 'PQAvgPool1d' object has no attribute 'pooling' occurs in many tests
• In the pytorch tests, they all fail with RuntimeError: Can't call numpy() on Tensor that requires grad. Use tensor.detach().numpy() instead.

[JanFSchulte]

Thanks! I think you need to catch up to change in the keras3 parser where the register decorator has been removed.

[JanFSchulte]

I don't quite get the logic in this mapping, can you please explain it?

[enlupi]

Let's assume we have two consecutive layers A, who has an input and output quantizer, and B.
Layer A will be parsed as 3 layers: FixedPointQuantizer A_iq, A proper, and FixedPointQuantizer A_oq. The input argument for the node B, though, will still be A, as A_oq did not exist at the moment of tracing the model. This means that the correct flow of data in the model will not be reconstructed, and both A_oq and B will be seen as outputs of A (A_oq will also be one of the outputs of the global model, as it is not used as input by anyone).

What these lines do is ensuring that if a layer sees A as input, it gets redirected to the next parsed layer, i.e. A_oq, and the correct flow can be reconstructed.

Admittedly, this has the useless consequence of redirecting the layers with A_iq as node argument to A, which should never happen as A_iq did not exist at tracing time.
Perhaps it could be reworked to just work on the A equivalent with something along the lines of:

if lay['name'] == layer_name and idx != len(layer) - 1:
inputs_map[lay['name']] = inputs_map.get(layer[-1]['name'], layer[-1]['name'])

Alternatively, we could do an explicit check for PQuantML layers and set the input map before all of this, as it is done in line 209 for skip layers.

[JanFSchulte]

Thanks for the explanation. I think the suggestion with the explicit check to only work on A might be a good idea to not have unintended consequences some time in the future, even if the current implementation is save now.

[enlupi]

Got it.
I changed it so that I explicitly check if it is a PQuantML layer and if it has an output quantizer; only in this case, the inputs_map is changed.

[JanFSchulte]

Can we document also what happens if this note is ignored? Will things break, will results be silently incorrect?

[enlupi]

I was following HGQ2's note, as the underlying bit exact pass that is invoked is the same. From my understanding, explicit layer precision is simply ignored (as per the Model-wise Precision Inference doc), while if you have both auto precision in the config and invoke the bit exact pass it will lead to undefined behavior.

[JanFSchulte]

Looks good to me now, thanks! Lets run the test one last time and merge once they are done.