Fix NaN/Infinity serialization in save_to_json_file (#4144)#5251
Fix NaN/Infinity serialization in save_to_json_file (#4144)#5251SiddharthSingh9018 wants to merge 2 commits into
Conversation
|
Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
saitcakmak
left a comment
There was a problem hiding this comment.
With this, we're intro
| try: | ||
| value = json["value"] | ||
| if isinstance(value, list): | ||
| from ax.storage.json_store.decoder import object_from_json |
There was a problem hiding this comment.
I'd like to avoid this circular import. For other cases like numpy arrays, we often decode the underlying object in object_from_json before calling the decoder from another file (like tensor_from_json). Let's repeat this pattern and instead call object_json["value"] = _object_from_json(object_json["value"]) before passing it down here. That will also generalize this handling to 0-dim tensors. Those are currently not handled here due to isinstance(value, list) check.
| elif _type is np.ndarray or issubclass(_type, np.ndarray): | ||
| return {"__type": _type.__name__, "value": obj.tolist()} | ||
| return {"__type": _type.__name__, "value": _object_to_json(obj.tolist())} | ||
| elif _type is set: | ||
| return {"__type": _type.__name__, "value": list(obj)} | ||
| return {"__type": _type.__name__, "value": _object_to_json(list(obj))} | ||
| elif _type is torch.Tensor: | ||
| return tensor_to_dict(obj=obj) | ||
| return {k: _object_to_json(v) for k, v in tensor_to_dict(obj=obj).items()} |
There was a problem hiding this comment.
This is introducing a python loop over each element of arrays / tensors. It could be worth gating the _object_to_json calls behind checks for inf / nan elements to avoid paying the cost for all objects
There was a problem hiding this comment.
Thanks,
For the circular import, I'll follow the existing pattern and move the recursive decode of object_json["value"] into object_from_json in decoder.py, before dispatching to decoders.py.
That also fixes the 0-dim tensor case you flagged, since the recursion won't depend on isinstance(value, list) anymore.
Plan:
Update decoder.py / decoders.py per the circular import fix, remove the local import
Add a test for a 0-dim tensor containing NaN
For the encoder perf concern: add a check that scans the array first to see if it actually contains any NaN/Inf
If all values are normal numbers (common case) → skip per-element conversion, use the existing fast path
If NaN/Inf is present (rare case) → fall back to per-element conversion
Skip the check entirely for non-float dtypes, since NaN/Inf doesn't apply there
Leave sets on a simple loop, since there's no vectorized check available
|
Updated in 67948da, addressing both review comments:
|
|
@saitcakmak has imported this pull request. If you are a Meta employee, you can view this in D112551642. |
Fixes #4144
Summary
This makes Ax JSON serialization standards-compliant when non-finite floats appear in the serialized object graph.
nan,inf, and-infusing Ax's typed JSON format.nan,inf, and-inf.Client.save_to_json_file()calljson.dumps(..., allow_nan=False)after conversion.Validation
python -m ufmt format ax\api\client.py ax\api\tests\test_client.py ax\storage\json_store\encoder.py ax\storage\json_store\decoder.py ax\storage\json_store\decoders.py ax\storage\json_store\tests\test_json_store.pypython -m py_compile ax\api\client.py ax\api\tests\test_client.py ax\storage\json_store\encoder.py ax\storage\json_store\decoder.py ax\storage\json_store\decoders.py ax\storage\json_store\tests\test_json_store.pyNote: focused pytest collection on Windows is blocked by Ax's
TestCaseusingsignal.SIGALRM, which Windows does not support.