Issue 1:
I'm seeing an excessive loss of precision far beyond what is typically acceptable for 64-bit calculations:

from jax.config import config; config.update("jax_enable_x64", True)

import jax
import jax.numpy as np
import numpy as anp
from jax.test_util import check_grads
import unittest

class PeriodicTorsion():

def __init__(self,
torsion_idxs,
param_idxs):
"""
This implements a periodic torsional potential expanded out into three terms:

V(a) = k0*(1+cos(1 * a - t0)) + k1*(1+cos(2 * a - t1)) + k2*(1+cos(3 * a - t2))

Parameters:
-----------
torsion_idxs: [num_torsions, 4] np.array
each element (a, b, c, d) is a torsion of four atoms, defined as
as the angle of the plane defined by the three bond vectors a-b, b-c, c-d.

param_idxs: [num_torsions, 6] np.array
each element (k, phase, periodicity) maps into params for angle constants and ideal angles

"""
self.torsion_idxs = torsion_idxs
self.param_idxs = param_idxs

@staticmethod
def get_signed_angle(ci, cj, ck, cl):
"""
The torsion angle between two planes should be periodic but not
necessarily symmetric. We use an identical but numerically stable arctan2
implementation as opposed to the OpenMM energy function to avoid a
singularity when the angle is zero. This is different from the CHARMM convention
of symmetric torsions.
"""

# Taken from the wikipedia arctan2 implementation:
# https://en.wikipedia.org/wiki/Dihedral_angle

rij = ci - cj
rkj = ck - cj
rkl = ck - cl

n1 = np.cross(rij, rkj)
n2 = np.cross(rkj, rkl)

lhs = np.linalg.norm(n1, axis=-1)
rhs = np.linalg.norm(n2, axis=-1)
bot = lhs * rhs

y = np.sum(np.multiply(np.cross(n1, n2), rkj/np.linalg.norm(rkj, axis=-1, keepdims=True)), axis=-1)
x = np.sum(np.multiply(n1, n2), -1)

return np.arctan2(y, x)

def angles(self, conf):
ci = conf[self.torsion_idxs[:, 0]]
cj = conf[self.torsion_idxs[:, 1]]
ck = conf[self.torsion_idxs[:, 2]]
cl = conf[self.torsion_idxs[:, 3]]

angle = self.get_signed_angle(ci, cj, ck, cl)
return angle

def energy(self, conf, params):
"""
Compute the torsional energy.
"""
ci = conf[self.torsion_idxs[:, 0]]
cj = conf[self.torsion_idxs[:, 1]]
ck = conf[self.torsion_idxs[:, 2]]
cl = conf[self.torsion_idxs[:, 3]]

ks = params[self.param_idxs[:, 0]]
phase = params[self.param_idxs[:, 1]]
period = params[self.param_idxs[:, 2]]
angle = self.get_signed_angle(ci, cj, ck, cl)
nrg = ks*(1+np.cos(period * angle - phase))
return np.sum(nrg, axis=-1)

class TestPeriodicTorsion(unittest.TestCase):

def setUp(self):
self.conformers = anp.array([
[[-0.6000563454193615, 0.376172954382274 ,-0.2487295756125901],
[ 0.561317027011325 , 0.2066950040043141, 0.3670430960815993],
[-1.187055522272264 ,-0.3415864358441354, 0.0871382207830652],
[ 0.9399773448903637,-0.6888774474110431, 0.2104211949995816]],
[[-0.6000563454193615, 0.376172954382274 ,-0.2487295756125901],
[ 0.5613170270113252, 0.2066950040043142, 0.3670430960815993],
[-1.187055522272264 ,-0.3415864358441354, 0.0871382207830652],
[ 1.283345455745044 ,-0.0356257425880843,-0.2573923896494185]],
[[-0.6000563454193615, 0.376172954382274 ,-0.2487295756125901],
[ 0.561317027011325 , 0.2066950040043142, 0.3670430960815992],
[-1.187055522272264 ,-0.3415864358441354, 0.0871382207830652],
[ 1.263820400176392 , 0.7964992122869241, 0.0084568741589791]],
[[-0.6000563454193615, 0.376172954382274 ,-0.2487295756125901],
[ 0.5613170270113252, 0.2066950040043142, 0.3670430960815992],
[-1.187055522272264 ,-0.3415864358441354, 0.0871382207830652],
[ 0.8993534242298198, 1.042445571242743 , 0.7635483993060286]],
[[-0.6000563454193615, 0.376172954382274 ,-0.2487295756125901],
[ 0.5613170270113255, 0.2066950040043142, 0.3670430960815993],
[-1.187055522272264 ,-0.3415864358441354, 0.0871382207830652],
[ 0.5250337847650304, 0.476091386095139 , 1.3136545198545133]],
[[-0.6000563454193615, 0.376172954382274 ,-0.2487295756125901],
[ 0.5613170270113255, 0.2066950040043141, 0.3670430960815993],
[-1.187055522272264 ,-0.3415864358441354, 0.0871382207830652],
[ 0.485009232042489 ,-0.3818599172073237, 1.1530102055165103]],
], dtype=anp.float64)

self.nan_conformers = anp.array([
[[-0.6000563454193615, 0.376172954382274 ,-0.2487295756125901],
[ 0.5613170270113252, 0.2066950040043142, 0.3670430960815993],
[-1.187055522272264 ,-0.3415864358441354, 0.0871382207830652],
[ 1.2278668040866427, 0.8805184219394547, 0.099391329616366 ]],
[[-0.6000563454193615, 0.376172954382274 ,-0.2487295756125901],
[ 0.561317027011325 , 0.206695004004314 , 0.3670430960815994],
[-1.187055522272264 ,-0.3415864358441354, 0.0871382207830652],
[ 0.5494071252089705,-0.5626592973923106, 0.9817919758125693]],
], dtype=anp.float64)

def test_gpu_torsions(self):
"""
Test agreement of torsions with OpenMM's implementation of torsion terms.
"""
torsion_idxs = anp.array([
[0, 1, 2, 3],
# [0, 1, 2, 3],
# [0, 1, 2, 3],
], dtype=anp.int32)

params_np = anp.array([
2.3, # k0
5.4, # k1
9.0, # k2
0.0, # t0
3.0, # t1
5.8, # t2
1.0, # n0
2.0, # n1
3.0 # n2
], dtype=anp.float64)

param_idxs = anp.array([
# [0, 3, 6], # works
[1, 4, 7], # fails
# [2, 5, 8] # fails
], dtype=anp.int32)

ref_nrg = PeriodicTorsion(
param_idxs=param_idxs,
torsion_idxs=torsion_idxs
)

for conf_idx, conf in enumerate(self.conformers):
nrg = ref_nrg.energy(conf, params_np)
angles = ref_nrg.angles(conf)

print(nrg, angles)

# first order derivatives pass
check_grads(ref_nrg.angles, (conf,), order=1) # passes
check_grads(ref_nrg.energy, (conf, params_np), order=1) # passes

# second order derivatives fail miserably, both of these should pass in 64bit
check_grads(ref_nrg.angles, (conf,), order=2) # fails
check_grads(ref_nrg.energy, (conf, params_np), order=2) # fails

if __name__ == "__main__":
unittest.main()

The failures look like:

======================================================================
FAIL: test_gpu_torsions (__main__.TestPeriodicTorsion)
----------------------------------------------------------------------
Traceback (most recent call last):
File "timemachine/jax_functionals/loss_of_precision_repro.py", line 175, in test_gpu_torsions
check_grads(ref_nrg.energy, (conf, params_np), order=2) # fails
File "/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py", line 153, in check_grads
check_grads(f_vjp, args, order - 1, atol=atol, rtol=rtol, eps=eps)
File "/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py", line 160, in check_grads
check_vjp(f, partial(api.vjp, f), args, atol, rtol, eps)
File "/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py", line 144, in check_vjp
check_close(ip, ip_expected, atol=atol, rtol=rtol)
File "/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py", line 89, in check_close
assert tree_all(tree_multimap(close, xs, ys)), '\n{} != \n{}'.format(xs, ys)
AssertionError:
55.438952303284566 !=
55.43906353281106

======================================================================
FAIL: test_gpu_torsions (__main__.TestPeriodicTorsion)
----------------------------------------------------------------------
Traceback (most recent call last):
File "timemachine/jax_functionals/loss_of_precision_repro.py", line 174, in test_gpu_torsions
check_grads(ref_nrg.angles, (conf,), order=2) # fails
File "/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py", line 153, in check_grads
check_grads(f_vjp, args, order - 1, atol=atol, rtol=rtol, eps=eps)
File "/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py", line 160, in check_grads
check_vjp(f, partial(api.vjp, f), args, atol, rtol, eps)
File "/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py", line 144, in check_vjp
check_close(ip, ip_expected, atol=atol, rtol=rtol)
File "/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py", line 89, in check_close
assert tree_all(tree_multimap(close, xs, ys)), '\n{} != \n{}'.format(xs, ys)
AssertionError:
0.381797537133656 !=
0.381794076418857

Issue 2:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

If applicable, include full error messages/tracebacks.

I am trying to optimize two parameters for linear regression. This is happening when a large number of points are involved.

I made a loop, but it failed silently. When I tried to see the updated parameters, I saw two nans.

import jax.numpy as jnp
import jax
import jaxlib

x_ = jnp.arange(2e7)
y_ = jnp.arange(2e7, 4e7)

theta = jnp.array([1., 1])

@jax.jit
def model(theta: jaxlib.xla_extension.DeviceArray, x: jaxlib.xla_extension.DeviceArray) -> jaxlib.xla_extension.DeviceArray:
w, b = theta
return w * x + b

@jax.jit
def loss_fn(theta: jaxlib.xla_extension.DeviceArray,
x: jaxlib.xla_extension.DeviceArray,
y: jaxlib.xla_extension.DeviceArray) -> jaxlib.xla_extension.DeviceArray:
'''
Calculates the MSE loss.
'''
prediction = model(theta, x) # y_hat
return jnp.mean((prediction - y) ** 2) # J

@jax.jit
def update(theta: jaxlib.xla_extension.DeviceArray,
x: jaxlib.xla_extension.DeviceArray,
y: jaxlib.xla_extension.DeviceArray,
lr: float = 0.1) -> jaxlib.xla_extension.DeviceArray:
'''
updates the parameters based on gradient descent
'''
return theta - lr * jax.grad(loss_fn)(theta, x, y)

for _ in range(1000):
theta = update(theta, x_, y_)

This failed silently.

Then I followed the docs and enables nan-parsing by-

!JAX_DEBUG_NANS=True

And then I did-

from jax.config import config
config.update("jax_debug_nans", True)

Which threw me a error.

After the error was thrown, I inspected the theta variable, and saw that it had very small values, i.e. the values before they were going nan.

It was this:

>>> theta
DeviceArray([-1.6214236e+37, -7.4782789e+35], dtype=float32)

What is going wrong? Am I missing something?

Note that I also tried initially with small values for the arrays. All values in array were temperature values, and the data is good (no outliers or wrong entry that might throw the parameters off).

Full error message:

Invalid nan value encountered in the output of a C++-jit function. Calling the de-optimized version.
Invalid value encountered in the output of a jit function. Calling the de-optimized version.
Invalid value encountered in the output of a jit function. Calling the de-optimized version.
Invalid value encountered in the output of a jit function. Calling the de-optimized version.
Invalid value encountered in the output of a jit function. Calling the de-optimized version.

---------------------------------------------------------------------------
FloatingPointError Traceback (most recent call last)
/opt/conda/lib/python3.7/site-packages/jax/_src/api.py in _nan_check_posthook(fun, args, kwargs, output)
120 da.device_buffer
--> 121 for da in tree_leaves(output)
122 if hasattr(da, "device_buffer")

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in check_special(name, bufs)
409 for buf in bufs:
--> 410 _check_special(name, buf.xla_shape(), buf)
411

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in _check_special(name, xla_shape, buf)
415 if config.jax_debug_nans and np.any(np.isnan(buf.to_py())):
--> 416 raise FloatingPointError(f"invalid value (nan) encountered in {name}")
417 if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):

FloatingPointError: invalid value (nan) encountered in xla_call

During handling of the above exception, another exception occurred:

FloatingPointError Traceback (most recent call last)
[... skipping hidden 1 frame]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in _execute_compiled(compiled, avals, handlers, kept_var_idx, *args)
912 out_bufs = compiled.execute(input_bufs)
--> 913 check_special(xla_call_p.name, out_bufs)
914 return [handler(*bs) for handler, bs in zip(handlers, _partition_outputs(avals, out_bufs))]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in check_special(name, bufs)
409 for buf in bufs:
--> 410 _check_special(name, buf.xla_shape(), buf)
411

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in _check_special(name, xla_shape, buf)
415 if config.jax_debug_nans and np.any(np.isnan(buf.to_py())):
--> 416 raise FloatingPointError(f"invalid value (nan) encountered in {name}")
417 if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):

FloatingPointError: invalid value (nan) encountered in xla_call

During handling of the above exception, another exception occurred:

FloatingPointError Traceback (most recent call last)
[... skipping hidden 1 frame]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in _execute_compiled(compiled, avals, handlers, kept_var_idx, *args)
912 out_bufs = compiled.execute(input_bufs)
--> 913 check_special(xla_call_p.name, out_bufs)
914 return [handler(*bs) for handler, bs in zip(handlers, _partition_outputs(avals, out_bufs))]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in check_special(name, bufs)
409 for buf in bufs:
--> 410 _check_special(name, buf.xla_shape(), buf)
411

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in _check_special(name, xla_shape, buf)
415 if config.jax_debug_nans and np.any(np.isnan(buf.to_py())):
--> 416 raise FloatingPointError(f"invalid value (nan) encountered in {name}")
417 if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):

FloatingPointError: invalid value (nan) encountered in xla_call

During handling of the above exception, another exception occurred:

FloatingPointError Traceback (most recent call last)
[... skipping hidden 1 frame]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in _execute_compiled(compiled, avals, handlers, kept_var_idx, *args)
912 out_bufs = compiled.execute(input_bufs)
--> 913 check_special(xla_call_p.name, out_bufs)
914 return [handler(*bs) for handler, bs in zip(handlers, _partition_outputs(avals, out_bufs))]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in check_special(name, bufs)
409 for buf in bufs:
--> 410 _check_special(name, buf.xla_shape(), buf)
411

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in _check_special(name, xla_shape, buf)
415 if config.jax_debug_nans and np.any(np.isnan(buf.to_py())):
--> 416 raise FloatingPointError(f"invalid value (nan) encountered in {name}")
417 if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):

FloatingPointError: invalid value (nan) encountered in xla_call

During handling of the above exception, another exception occurred:

FloatingPointError Traceback (most recent call last)
[... skipping hidden 1 frame]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in _execute_compiled(compiled, avals, handlers, kept_var_idx, *args)
912 out_bufs = compiled.execute(input_bufs)
--> 913 check_special(xla_call_p.name, out_bufs)
914 return [handler(*bs) for handler, bs in zip(handlers, _partition_outputs(avals, out_bufs))]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in check_special(name, bufs)
409 for buf in bufs:
--> 410 _check_special(name, buf.xla_shape(), buf)
411

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in _check_special(name, xla_shape, buf)
415 if config.jax_debug_nans and np.any(np.isnan(buf.to_py())):
--> 416 raise FloatingPointError(f"invalid value (nan) encountered in {name}")
417 if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):

FloatingPointError: invalid value (nan) encountered in xla_call

During handling of the above exception, another exception occurred:

JaxStackTraceBeforeTransformation Traceback (most recent call last)
/opt/conda/lib/python3.7/runpy.py in _run_module_as_main(***failed resolving arguments***)
192 return _run_code(code, main_globals, None,
--> 193 "__main__", mod_spec)
194

/opt/conda/lib/python3.7/runpy.py in _run_code(***failed resolving arguments***)
84 __spec__ = mod_spec)
---> 85 exec(code, run_globals)
86 return run_globals

/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py in <module>
15 from ipykernel import kernelapp as app
---> 16 app.launch_new_instance()

/opt/conda/lib/python3.7/site-packages/traitlets/config/application.py in launch_instance(***failed resolving arguments***)
844 app.initialize(argv)
--> 845 app.start()
846

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelapp.py in start(***failed resolving arguments***)
666 try:
--> 667 self.io_loop.start()
668 except KeyboardInterrupt:

/opt/conda/lib/python3.7/site-packages/tornado/platform/asyncio.py in start(***failed resolving arguments***)
198 asyncio.set_event_loop(self.asyncio_loop)
--> 199 self.asyncio_loop.run_forever()
200 finally:

/opt/conda/lib/python3.7/asyncio/base_events.py in run_forever(***failed resolving arguments***)
540 while True:
--> 541 self._run_once()
542 if self._stopping:

/opt/conda/lib/python3.7/asyncio/base_events.py in _run_once(***failed resolving arguments***)
1785 else:
-> 1786 handle._run()
1787 handle = None # Needed to break cycles when an exception occurs.

/opt/conda/lib/python3.7/asyncio/events.py in _run(***failed resolving arguments***)
87 try:
---> 88 self._context.run(self._callback, *self._args)
89 except Exception as exc:

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py in dispatch_queue(***failed resolving arguments***)
456 try:
--> 457 await self.process_one()
458 except Exception:

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py in process_one(***failed resolving arguments***)
445 return None
--> 446 await dispatch(*args)
447

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py in dispatch_shell(***failed resolving arguments***)
352 if inspect.isawaitable(result):
--> 353 await result
354 except Exception:

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py in execute_request(***failed resolving arguments***)
647 if inspect.isawaitable(reply_content):
--> 648 reply_content = await reply_content
649

/opt/conda/lib/python3.7/site-packages/ipykernel/ipkernel.py in do_execute(***failed resolving arguments***)
344 # letting shell dispatch to loop runners
--> 345 res = shell.run_cell(code, store_history=store_history, silent=silent)
346 finally:

/opt/conda/lib/python3.7/site-packages/ipykernel/zmqshell.py in run_cell(***failed resolving arguments***)
531 self._last_traceback = None
--> 532 return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
533

/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py in run_cell(***failed resolving arguments***)
2898 result = self._run_cell(
-> 2899 raw_cell, store_history, silent, shell_futures)
2900 finally:

/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py in _run_cell(***failed resolving arguments***)
2943 try:
-> 2944 return runner(coro)
2945 except BaseException as e:

/opt/conda/lib/python3.7/site-packages/IPython/core/async_helpers.py in _pseudo_sync_runner(***failed resolving arguments***)
67 try:
---> 68 coro.send(None)
69 except StopIteration as exc:

/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py in run_cell_async(***failed resolving arguments***)
3169 has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
-> 3170 interactivity=interactivity, compiler=compiler, result=result)
3171

/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py in run_ast_nodes(***failed resolving arguments***)
3360 asy = compare(code)
-> 3361 if (await self.run_code(code, result, async_=asy)):
3362 return True

/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py in run_code(***failed resolving arguments***)
3440 else:
-> 3441 exec(code_obj, self.user_global_ns, self.user_ns)
3442 finally:

/tmp/ipykernel_43/2138327653.py in <module>
1 for _ in range(1000):
----> 2 theta = update(theta, x_train, y_train)
3

/tmp/ipykernel_43/1444374627.py in update(***failed resolving arguments***)
8 '''
----> 9 return theta - lr * jax.grad(loss_fn)(theta, x, y)

/tmp/ipykernel_43/3487162561.py in loss_fn(***failed resolving arguments***)
7 '''
----> 8 prediction = model(theta, x) # y_hat
9 return jnp.mean((prediction - y) ** 2) # J

/tmp/ipykernel_43/3906416787.py in model(***failed resolving arguments***)
3 w, b = theta
----> 4 return w * x + b

/opt/conda/lib/python3.7/site-packages/jax/_src/numpy/lax_numpy.py in deferring_binary_op(***failed resolving arguments***)
5923 return NotImplemented
-> 5924 return binary_op(self, other)
5925 return deferring_binary_op

/opt/conda/lib/python3.7/site-packages/jax/_src/numpy/lax_numpy.py in fn(***failed resolving arguments***)
435 x1, x2 = _promote_args(numpy_fn.__name__, x1, x2)
--> 436 return lax_fn(x1, x2) if x1.dtype != bool_ else bool_lax_fn(x1, x2)
437 fn = jit(fn, inline=True)

JaxStackTraceBeforeTransformation: FloatingPointError: invalid value (nan) encountered in reduce_sum

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

FloatingPointError Traceback (most recent call last)
/tmp/ipykernel_43/2138327653.py in <module>
1 for _ in range(1000):
----> 2 theta = update(theta, x_train, y_train)
3
4 plt.scatter(x_train, y_train)
5 plt.plot(x_train, model(theta, x_train), 'r')

/opt/conda/lib/python3.7/site-packages/jax/_src/api.py in _nan_check_posthook(fun, args, kwargs, output)
127 print("Invalid nan value encountered in the output of a C++-jit "
128 "function. Calling the de-optimized version.")
--> 129 fun._cache_miss(*args, **kwargs)[0] # probably won't return
130
131 def _update_debug_special_global(_):

[... skipping hidden 8 frame]

/tmp/ipykernel_43/1444374627.py in update(theta, x, y, lr)
7 updates the parameters based on gradient descent
8 '''
----> 9 return theta - lr * jax.grad(loss_fn)(theta, x, y)

[... skipping hidden 39 frame]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in _check_special(name, xla_shape, buf)
414 if dtypes.issubdtype(xla_shape.element_type(), np.inexact):
415 if config.jax_debug_nans and np.any(np.isnan(buf.to_py())):
--> 416 raise FloatingPointError(f"invalid value (nan) encountered in {name}")
417 if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):
418 raise FloatingPointError(f"invalid value (inf) encountered in {name}")

FloatingPointError: invalid value (nan) encountered in reduce_sum