fix: Correctly handle inf and nan values

Correctly handle inf and nan values. Fall back on the standard library's
`json` module in the rare event that JSON lines contain nan or inf
values.

stan/fit.py

@@ -1,3 +1,4 @@
+import json
 from typing import Sequence
 
 import numpy as np
@@ -64,10 +65,14 @@ def __init__(
         for chain_index, stan_output in zip(range(self.num_chains), self.stan_outputs):
             draw_index = 0
             for line in stan_output.splitlines():
-                msg = parser.parse(line)
+                try:
+                    msg = parser.parse(line)
+                except ValueError:
+                    # Occurs when draws contain an nan or infinity. simdjson cannot parse such values.
+                    msg = json.loads(line)
                 if msg["topic"] == "sample":
                     # Ignore sample message which is mixed together with proper draws.
-                    if not isinstance(msg["values"], simdjson.Object):
+                    if not isinstance(msg["values"], (simdjson.Object, dict)):
                         continue
 
                     # for the first draw: collect sample and sampler parameter names.

stan/model.py

@@ -195,6 +195,10 @@ def is_iteration_or_elapsed_time_logger_message(msg: simdjson.Object):
                 for stan_output in stan_outputs:
                     assert isinstance(stan_output, bytes)
                     for line in stan_output.splitlines():
+                        # Do not attempt to parse non-logger messages. Draws could contain nan or inf values.
+                        # simdjson cannot parse lines containing such values.
+                        if b'"logger"' not in line:
+                            continue
                         msg = parser.parse(line)
                         if is_nonempty_logger_message(msg) and not is_iteration_or_elapsed_time_logger_message(msg):
                             nonstandard_logger_messages.append(msg.as_dict())