[JAX] Align and extend docs and examples - #2567
Open
bkowalskiINTEL wants to merge 3 commits into
Open
Conversation
…and example Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>
bkowalskiINTEL
force-pushed
the
dev/bkowalsk/jax_align_docs_and_examples
branch
from
September 4, 2026 02:09
bfd9dec to
6389b7d
Compare
…on and examples Signed-off-by: Bartosz Kowalski <bartosz.kowalski@intel.com>
anko-intel
reviewed
Sep 4, 2026
| ```python | ||
| def quantize_model( | ||
| model: keras.Model, | ||
| quant_config: BaseConfig, |
Contributor
There was a problem hiding this comment.
Suggested change
| quant_config: JaxBaseConfig, |
anko-intel
reviewed
Sep 4, 2026
| import numpy as np | ||
| import tensorflow as tf | ||
| from keras.applications.imagenet_utils import decode_predictions | ||
| from keras_hub.models import ViTImageClassifier |
Contributor
There was a problem hiding this comment.
Maybe better will be to move import ViTImageClassifier to line 40 to show that this import is only required when from_preset is used (?)
|
|
||
| from neural_compressor.jax import quantize_model, StaticQuantConfig | ||
|
|
||
| from keras_hub.models import ViTImageClassifier |
Contributor
There was a problem hiding this comment.
maybe this import only for line 46 and 76 (?)
| print(f"\nLoading quantization configuration from: {args.quant_config_file}") | ||
| config = JaxBaseConfig.from_json_file(args.quant_config_file) | ||
| else: | ||
| print("\nUsing default composable quantization configuration") |
Contributor
There was a problem hiding this comment.
Suggested change
| print("\nUsing default composable quantization configuration") | |
| print("\nPreparing composable quantization configuration") |
| } | ||
| } | ||
| ] | ||
| } No newline at end of file |
Contributor
There was a problem hiding this comment.
Add just new line at the and to avoid GitHub "-" mark
| - [Introduction](#introduction) | ||
| - [Quantization API](#quantization-api) | ||
| - [Post-Training Static Quantization](#post-training-static-quantization) | ||
| - [Examples](#examples) |
| """ | ||
| ``` | ||
|
|
||
| ## Quantization configs |
Contributor
There was a problem hiding this comment.
I am not sure if detailed description of configs should be in this main README. Maybe it could be in config examples
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Type of Change
documentation
Description
This PR adds more detailed documentation about config serialization and composable configs, and provides config json file examples so that the users can have better understanding of how to create custom ones.
It also adds a small change in vit examples enabling both save_model/load_model and save_to_preset/from_preset APIs.