Skip to content

Fix NaN/Infinity serialization in save_to_json_file (#4144)#5251

Open
SiddharthSingh9018 wants to merge 2 commits into
facebook:mainfrom
SiddharthSingh9018:fix-nan-json-serialization
Open

Fix NaN/Infinity serialization in save_to_json_file (#4144)#5251
SiddharthSingh9018 wants to merge 2 commits into
facebook:mainfrom
SiddharthSingh9018:fix-nan-json-serialization

Conversation

@SiddharthSingh9018

Copy link
Copy Markdown

Fixes #4144

Summary

This makes Ax JSON serialization standards-compliant when non-finite floats appear in the serialized object graph.

  • Encodes nan, inf, and -inf using Ax's typed JSON format.
  • Handles Python floats and NumPy floating values.
  • Recursively handles nested containers, ndarray payloads, tensor payloads, and sets.
  • Decodes typed float values back to nan, inf, and -inf.
  • Makes Client.save_to_json_file() call json.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.py
  • python -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.py
  • Direct strict JSON save/load verification passed.

Note: focused pytest collection on Windows is blocked by Ax's TestCase using signal.SIGALRM, which Windows does not support.

@meta-cla

meta-cla Bot commented Jul 9, 2026

Copy link
Copy Markdown

Hi @SiddharthSingh9018!

Thank you for your pull request and welcome to our community.

Action Required

In 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.

Process

In 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 CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

@meta-cla meta-cla Bot added the CLA Signed Do not delete this pull request or issue due to inactivity. label Jul 9, 2026

@saitcakmak saitcakmak left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With this, we're intro

Comment thread ax/storage/json_store/decoders.py Outdated
try:
value = json["value"]
if isinstance(value, list):
from ax.storage.json_store.decoder import object_from_json

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread ax/storage/json_store/encoder.py Outdated
Comment on lines +119 to +124
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()}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@SiddharthSingh9018

Copy link
Copy Markdown
Author

Updated in 67948da, addressing both review comments:

  • moved recursive value decoding into object_from_json to remove the tensor decoder circular import and cover 0-dimensional tensor/array values
  • kept finite ndarray/tensor/set serialization on the fast path while only falling back to per-element conversion for non-finite values
    I also added regression coverage for 0-dimensional non-finite ndarray/tensor values and finite ndarray/tensor output stability.

@saitcakmak saitcakmak left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@meta-codesync

meta-codesync Bot commented Jul 17, 2026

Copy link
Copy Markdown

@saitcakmak has imported this pull request. If you are a Meta employee, you can view this in D112551642.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed Do not delete this pull request or issue due to inactivity.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: save_to_json_file saves a value NaN which not compatible with json format

2 participants