Skip to content

Add initial structure for HiJAX integration in qwix.#249

Open
copybara-service[bot] wants to merge 1 commit intomainfrom
test_897887809
Open

Add initial structure for HiJAX integration in qwix.#249
copybara-service[bot] wants to merge 1 commit intomainfrom
test_897887809

Conversation

@copybara-service
Copy link
Copy Markdown

Add initial structure for HiJAX integration in qwix.

This change introduces the qwix/contrib/hijax directory, which will house an experimental implementation of QArrays using HiJAX. The README explains the motivation for using HiJAX to gain finer control at the jaxpr level.

The CL includes:

  • A new BUILD file.
  • hiqarray_common.py: Defines QuantizationMetadata to manage shapes and axes for quantization.
  • hiquant_utils.py: Provides utility functions for handling dtypes in quantization.
  • hiqarray_common_test.py: Tests for the shape and axis computation in QuantizationMetadata.
  • A detailed README.md file with motivation and implementation notes.

This change introduces the `qwix/contrib/hijax` directory, which will house an experimental implementation of QArrays using HiJAX. The README explains the motivation for using HiJAX to gain finer control at the jaxpr level.

The CL includes:
-   A new BUILD file.
-   `hiqarray_common.py`: Defines `QuantizationMetadata` to manage shapes and axes for quantization.
-   `hiquant_utils.py`: Provides utility functions for handling dtypes in quantization.
-   `hiqarray_common_test.py`: Tests for the shape and axis computation in `QuantizationMetadata`.
-   A detailed README.md file with motivation and implementation notes.

PiperOrigin-RevId: 897887809
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant