<p dir="auto">I'm seeing an excessive loss of precision far beyond what is typically acceptable for 64-bit calculations:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from jax.config import config; config.update(&quot;jax_enable_x64&quot;, 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):
        &quot;&quot;&quot;
        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

        &quot;&quot;&quot;
        self.torsion_idxs = torsion_idxs
        self.param_idxs = param_idxs

    @staticmethod
    def get_signed_angle(ci, cj, ck, cl):
        &quot;&quot;&quot;
        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.
        &quot;&quot;&quot;

        # 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):
        &quot;&quot;&quot;
        Compute the torsional energy.
        &quot;&quot;&quot;
        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):
        &quot;&quot;&quot;
        Test agreement of torsions with OpenMM's implementation of torsion terms.
        &quot;&quot;&quot;
        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__ == &quot;__main__&quot;:
    unittest.main()"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">config</span> <span class="pl-k">import</span> <span class="pl-s1">config</span>; <span class="pl-s1">config</span>.<span class="pl-en">update</span>(<span class="pl-s">"jax_enable_x64"</span>, <span class="pl-c1">True</span>)

<span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>
<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">anp</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">test_util</span> <span class="pl-k">import</span> <span class="pl-s1">check_grads</span>
<span class="pl-k">import</span> <span class="pl-s1">unittest</span>

<span class="pl-k">class</span> <span class="pl-v">PeriodicTorsion</span>():

    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>,
        <span class="pl-s1">torsion_idxs</span>,
        <span class="pl-s1">param_idxs</span>):
        <span class="pl-s">"""</span>
<span class="pl-s">        This implements a periodic torsional potential expanded out into three terms:</span>
<span class="pl-s"></span>
<span class="pl-s">        V(a) = k0*(1+cos(1 * a - t0)) + k1*(1+cos(2 * a - t1)) + k2*(1+cos(3 * a - t2))</span>
<span class="pl-s"></span>
<span class="pl-s">        Parameters:</span>
<span class="pl-s">        -----------</span>
<span class="pl-s">        torsion_idxs: [num_torsions, 4] np.array</span>
<span class="pl-s">            each element (a, b, c, d) is a torsion of four atoms, defined as</span>
<span class="pl-s">            as the angle of the plane defined by the three bond vectors a-b, b-c, c-d. </span>
<span class="pl-s"></span>
<span class="pl-s">        param_idxs: [num_torsions, 6] np.array</span>
<span class="pl-s">            each element (k, phase, periodicity) maps into params for angle constants and ideal angles</span>
<span class="pl-s"></span>
<span class="pl-s">        """</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">torsion_idxs</span> <span class="pl-c1">=</span> <span class="pl-s1">torsion_idxs</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">param_idxs</span> <span class="pl-c1">=</span> <span class="pl-s1">param_idxs</span>

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

        <span class="pl-c"># Taken from the wikipedia arctan2 implementation:</span>
        <span class="pl-c"># https://en.wikipedia.org/wiki/Dihedral_angle</span>

        <span class="pl-s1">rij</span> <span class="pl-c1">=</span> <span class="pl-s1">ci</span> <span class="pl-c1">-</span> <span class="pl-s1">cj</span>
        <span class="pl-s1">rkj</span> <span class="pl-c1">=</span> <span class="pl-s1">ck</span> <span class="pl-c1">-</span> <span class="pl-s1">cj</span>
        <span class="pl-s1">rkl</span> <span class="pl-c1">=</span> <span class="pl-s1">ck</span> <span class="pl-c1">-</span> <span class="pl-s1">cl</span>

        <span class="pl-s1">n1</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">cross</span>(<span class="pl-s1">rij</span>, <span class="pl-s1">rkj</span>)
        <span class="pl-s1">n2</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">cross</span>(<span class="pl-s1">rkj</span>, <span class="pl-s1">rkl</span>)

        <span class="pl-s1">lhs</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">n1</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)
        <span class="pl-s1">rhs</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">n2</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)
        <span class="pl-s1">bot</span> <span class="pl-c1">=</span> <span class="pl-s1">lhs</span> <span class="pl-c1">*</span> <span class="pl-s1">rhs</span>

        <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">sum</span>(<span class="pl-s1">np</span>.<span class="pl-en">multiply</span>(<span class="pl-s1">np</span>.<span class="pl-en">cross</span>(<span class="pl-s1">n1</span>, <span class="pl-s1">n2</span>), <span class="pl-s1">rkj</span><span class="pl-c1">/</span><span class="pl-s1">np</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>(<span class="pl-s1">rkj</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-s1">keepdims</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)), <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">sum</span>(<span class="pl-s1">np</span>.<span class="pl-en">multiply</span>(<span class="pl-s1">n1</span>, <span class="pl-s1">n2</span>), <span class="pl-c1">-</span><span class="pl-c1">1</span>)

        <span class="pl-k">return</span> <span class="pl-s1">np</span>.<span class="pl-en">arctan2</span>(<span class="pl-s1">y</span>, <span class="pl-s1">x</span>)

    <span class="pl-k">def</span> <span class="pl-en">angles</span>(<span class="pl-s1">self</span>, <span class="pl-s1">conf</span>):
        <span class="pl-s1">ci</span> <span class="pl-c1">=</span> <span class="pl-s1">conf</span>[<span class="pl-s1">self</span>.<span class="pl-s1">torsion_idxs</span>[:, <span class="pl-c1">0</span>]]
        <span class="pl-s1">cj</span> <span class="pl-c1">=</span> <span class="pl-s1">conf</span>[<span class="pl-s1">self</span>.<span class="pl-s1">torsion_idxs</span>[:, <span class="pl-c1">1</span>]]
        <span class="pl-s1">ck</span> <span class="pl-c1">=</span> <span class="pl-s1">conf</span>[<span class="pl-s1">self</span>.<span class="pl-s1">torsion_idxs</span>[:, <span class="pl-c1">2</span>]]
        <span class="pl-s1">cl</span> <span class="pl-c1">=</span> <span class="pl-s1">conf</span>[<span class="pl-s1">self</span>.<span class="pl-s1">torsion_idxs</span>[:, <span class="pl-c1">3</span>]]

        <span class="pl-s1">angle</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">get_signed_angle</span>(<span class="pl-s1">ci</span>, <span class="pl-s1">cj</span>, <span class="pl-s1">ck</span>, <span class="pl-s1">cl</span>)
        <span class="pl-k">return</span> <span class="pl-s1">angle</span>

    <span class="pl-k">def</span> <span class="pl-en">energy</span>(<span class="pl-s1">self</span>, <span class="pl-s1">conf</span>, <span class="pl-s1">params</span>):
        <span class="pl-s">"""</span>
<span class="pl-s">        Compute the torsional energy.</span>
<span class="pl-s">        """</span>
        <span class="pl-s1">ci</span> <span class="pl-c1">=</span> <span class="pl-s1">conf</span>[<span class="pl-s1">self</span>.<span class="pl-s1">torsion_idxs</span>[:, <span class="pl-c1">0</span>]]
        <span class="pl-s1">cj</span> <span class="pl-c1">=</span> <span class="pl-s1">conf</span>[<span class="pl-s1">self</span>.<span class="pl-s1">torsion_idxs</span>[:, <span class="pl-c1">1</span>]]
        <span class="pl-s1">ck</span> <span class="pl-c1">=</span> <span class="pl-s1">conf</span>[<span class="pl-s1">self</span>.<span class="pl-s1">torsion_idxs</span>[:, <span class="pl-c1">2</span>]]
        <span class="pl-s1">cl</span> <span class="pl-c1">=</span> <span class="pl-s1">conf</span>[<span class="pl-s1">self</span>.<span class="pl-s1">torsion_idxs</span>[:, <span class="pl-c1">3</span>]]

        <span class="pl-s1">ks</span> <span class="pl-c1">=</span> <span class="pl-s1">params</span>[<span class="pl-s1">self</span>.<span class="pl-s1">param_idxs</span>[:, <span class="pl-c1">0</span>]]
        <span class="pl-s1">phase</span> <span class="pl-c1">=</span> <span class="pl-s1">params</span>[<span class="pl-s1">self</span>.<span class="pl-s1">param_idxs</span>[:, <span class="pl-c1">1</span>]]
        <span class="pl-s1">period</span> <span class="pl-c1">=</span> <span class="pl-s1">params</span>[<span class="pl-s1">self</span>.<span class="pl-s1">param_idxs</span>[:, <span class="pl-c1">2</span>]]
        <span class="pl-s1">angle</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">get_signed_angle</span>(<span class="pl-s1">ci</span>, <span class="pl-s1">cj</span>, <span class="pl-s1">ck</span>, <span class="pl-s1">cl</span>)
        <span class="pl-s1">nrg</span> <span class="pl-c1">=</span> <span class="pl-s1">ks</span><span class="pl-c1">*</span>(<span class="pl-c1">1</span><span class="pl-c1">+</span><span class="pl-s1">np</span>.<span class="pl-en">cos</span>(<span class="pl-s1">period</span> <span class="pl-c1">*</span> <span class="pl-s1">angle</span> <span class="pl-c1">-</span> <span class="pl-s1">phase</span>))
        <span class="pl-k">return</span> <span class="pl-s1">np</span>.<span class="pl-en">sum</span>(<span class="pl-s1">nrg</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)


<span class="pl-k">class</span> <span class="pl-v">TestPeriodicTorsion</span>(<span class="pl-s1">unittest</span>.<span class="pl-v">TestCase</span>):

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

        <span class="pl-s1">self</span>.<span class="pl-s1">nan_conformers</span> <span class="pl-c1">=</span> <span class="pl-s1">anp</span>.<span class="pl-en">array</span>([
            [[<span class="pl-c1">-</span><span class="pl-c1">0.6000563454193615</span>, <span class="pl-c1">0.376172954382274</span> ,<span class="pl-c1">-</span><span class="pl-c1">0.2487295756125901</span>],
             [ <span class="pl-c1">0.5613170270113252</span>, <span class="pl-c1">0.2066950040043142</span>, <span class="pl-c1">0.3670430960815993</span>],
             [<span class="pl-c1">-</span><span class="pl-c1">1.187055522272264</span> ,<span class="pl-c1">-</span><span class="pl-c1">0.3415864358441354</span>, <span class="pl-c1">0.0871382207830652</span>],
             [ <span class="pl-c1">1.2278668040866427</span>, <span class="pl-c1">0.8805184219394547</span>, <span class="pl-c1">0.099391329616366</span> ]],
            [[<span class="pl-c1">-</span><span class="pl-c1">0.6000563454193615</span>, <span class="pl-c1">0.376172954382274</span> ,<span class="pl-c1">-</span><span class="pl-c1">0.2487295756125901</span>],
             [ <span class="pl-c1">0.561317027011325</span> , <span class="pl-c1">0.206695004004314</span> , <span class="pl-c1">0.3670430960815994</span>],
             [<span class="pl-c1">-</span><span class="pl-c1">1.187055522272264</span> ,<span class="pl-c1">-</span><span class="pl-c1">0.3415864358441354</span>, <span class="pl-c1">0.0871382207830652</span>],
             [ <span class="pl-c1">0.5494071252089705</span>,<span class="pl-c1">-</span><span class="pl-c1">0.5626592973923106</span>, <span class="pl-c1">0.9817919758125693</span>]],
            ], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">anp</span>.<span class="pl-s1">float64</span>)


    <span class="pl-k">def</span> <span class="pl-en">test_gpu_torsions</span>(<span class="pl-s1">self</span>):
        <span class="pl-s">"""</span>
<span class="pl-s">        Test agreement of torsions with OpenMM's implementation of torsion terms.</span>
<span class="pl-s">        """</span>
        <span class="pl-s1">torsion_idxs</span> <span class="pl-c1">=</span> <span class="pl-s1">anp</span>.<span class="pl-en">array</span>([
            [<span class="pl-c1">0</span>, <span class="pl-c1">1</span>, <span class="pl-c1">2</span>, <span class="pl-c1">3</span>],
            <span class="pl-c"># [0, 1, 2, 3],</span>
            <span class="pl-c"># [0, 1, 2, 3],</span>
        ], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">anp</span>.<span class="pl-s1">int32</span>)

        <span class="pl-s1">params_np</span> <span class="pl-c1">=</span> <span class="pl-s1">anp</span>.<span class="pl-en">array</span>([
            <span class="pl-c1">2.3</span>, <span class="pl-c"># k0</span>
            <span class="pl-c1">5.4</span>, <span class="pl-c"># k1</span>
            <span class="pl-c1">9.0</span>, <span class="pl-c"># k2</span>
            <span class="pl-c1">0.0</span>, <span class="pl-c"># t0</span>
            <span class="pl-c1">3.0</span>, <span class="pl-c"># t1</span>
            <span class="pl-c1">5.8</span>, <span class="pl-c"># t2</span>
            <span class="pl-c1">1.0</span>, <span class="pl-c"># n0</span>
            <span class="pl-c1">2.0</span>, <span class="pl-c"># n1</span>
            <span class="pl-c1">3.0</span>  <span class="pl-c"># n2</span>
        ], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">anp</span>.<span class="pl-s1">float64</span>)

        <span class="pl-s1">param_idxs</span> <span class="pl-c1">=</span> <span class="pl-s1">anp</span>.<span class="pl-en">array</span>([
            <span class="pl-c"># [0, 3, 6], # works</span>
            [<span class="pl-c1">1</span>, <span class="pl-c1">4</span>, <span class="pl-c1">7</span>], <span class="pl-c"># fails</span>
            <span class="pl-c"># [2, 5, 8] # fails</span>
        ], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">anp</span>.<span class="pl-s1">int32</span>)

        <span class="pl-s1">ref_nrg</span> <span class="pl-c1">=</span> <span class="pl-v">PeriodicTorsion</span>(
            <span class="pl-s1">param_idxs</span><span class="pl-c1">=</span><span class="pl-s1">param_idxs</span>,
            <span class="pl-s1">torsion_idxs</span><span class="pl-c1">=</span><span class="pl-s1">torsion_idxs</span>
        )

        <span class="pl-k">for</span> <span class="pl-s1">conf_idx</span>, <span class="pl-s1">conf</span> <span class="pl-c1">in</span> <span class="pl-en">enumerate</span>(<span class="pl-s1">self</span>.<span class="pl-s1">conformers</span>):
            <span class="pl-s1">nrg</span> <span class="pl-c1">=</span> <span class="pl-s1">ref_nrg</span>.<span class="pl-en">energy</span>(<span class="pl-s1">conf</span>, <span class="pl-s1">params_np</span>)
            <span class="pl-s1">angles</span> <span class="pl-c1">=</span> <span class="pl-s1">ref_nrg</span>.<span class="pl-en">angles</span>(<span class="pl-s1">conf</span>)

            <span class="pl-en">print</span>(<span class="pl-s1">nrg</span>, <span class="pl-s1">angles</span>)

            <span class="pl-c"># first order derivatives pass</span>
            <span class="pl-en">check_grads</span>(<span class="pl-s1">ref_nrg</span>.<span class="pl-s1">angles</span>, (<span class="pl-s1">conf</span>,), <span class="pl-s1">order</span><span class="pl-c1">=</span><span class="pl-c1">1</span>) <span class="pl-c"># passes</span>
            <span class="pl-en">check_grads</span>(<span class="pl-s1">ref_nrg</span>.<span class="pl-s1">energy</span>, (<span class="pl-s1">conf</span>, <span class="pl-s1">params_np</span>), <span class="pl-s1">order</span><span class="pl-c1">=</span><span class="pl-c1">1</span>) <span class="pl-c"># passes</span>


            <span class="pl-c"># second order derivatives fail miserably, both of these should pass in 64bit</span>
            <span class="pl-en">check_grads</span>(<span class="pl-s1">ref_nrg</span>.<span class="pl-s1">angles</span>, (<span class="pl-s1">conf</span>,), <span class="pl-s1">order</span><span class="pl-c1">=</span><span class="pl-c1">2</span>) <span class="pl-c"># fails</span>
            <span class="pl-en">check_grads</span>(<span class="pl-s1">ref_nrg</span>.<span class="pl-s1">energy</span>, (<span class="pl-s1">conf</span>, <span class="pl-s1">params_np</span>), <span class="pl-s1">order</span><span class="pl-c1">=</span><span class="pl-c1">2</span>) <span class="pl-c"># fails</span>




<span class="pl-k">if</span> <span class="pl-s1">__name__</span> <span class="pl-c1">==</span> <span class="pl-s">"__main__"</span>:
    <span class="pl-s1">unittest</span>.<span class="pl-en">main</span>()</pre></div>
<p dir="auto">The failures look like:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="======================================================================
FAIL: test_gpu_torsions (__main__.TestPeriodicTorsion)
----------------------------------------------------------------------
Traceback (most recent call last):
  File &quot;timemachine/jax_functionals/loss_of_precision_repro.py&quot;, line 175, in test_gpu_torsions
    check_grads(ref_nrg.energy, (conf, params_np), order=2) # fails
  File &quot;/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py&quot;, line 153, in check_grads
    check_grads(f_vjp, args, order - 1, atol=atol, rtol=rtol, eps=eps)
  File &quot;/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py&quot;, line 160, in check_grads
    check_vjp(f, partial(api.vjp, f), args, atol, rtol, eps)
  File &quot;/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py&quot;, line 144, in check_vjp
    check_close(ip, ip_expected, atol=atol, rtol=rtol)
  File &quot;/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py&quot;, 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 &quot;timemachine/jax_functionals/loss_of_precision_repro.py&quot;, line 174, in test_gpu_torsions
    check_grads(ref_nrg.angles, (conf,), order=2) # fails
  File &quot;/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py&quot;, line 153, in check_grads
    check_grads(f_vjp, args, order - 1, atol=atol, rtol=rtol, eps=eps)
  File &quot;/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py&quot;, line 160, in check_grads
    check_vjp(f, partial(api.vjp, f), args, atol, rtol, eps)
  File &quot;/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py&quot;, line 144, in check_vjp
    check_close(ip, ip_expected, atol=atol, rtol=rtol)
  File &quot;/Users/hessian/venv/lib/python3.6/site-packages/jax/test_util.py&quot;, line 89, in check_close
    assert tree_all(tree_multimap(close, xs, ys)), '\n{} != \n{}'.format(xs, ys)
AssertionError: 
0.381797537133656 != 
0.381794076418857"><pre class="notranslate"><code class="notranslate">======================================================================
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
</code></pre></div>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<p dir="auto">I am trying to optimize two parameters for linear regression. This is happening when a large number of points are involved.</p>
<p dir="auto">I made a loop, but it failed silently. When I tried to see the updated parameters, I saw two <code class="notranslate">nan</code>s.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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) -&gt; 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) -&gt; 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) -&gt; 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_)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jaxlib</span>

<span class="pl-s1">x_</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">2e7</span>)
<span class="pl-s1">y_</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">2e7</span>, <span class="pl-c1">4e7</span>)

<span class="pl-s1">theta</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1.</span>, <span class="pl-c1">1</span>])

<span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">model</span>(<span class="pl-s1">theta</span>: <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>, <span class="pl-s1">x</span>: <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>) <span class="pl-c1">-&gt;</span> <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>:
    <span class="pl-s1">w</span>, <span class="pl-s1">b</span> <span class="pl-c1">=</span> <span class="pl-s1">theta</span>
    <span class="pl-k">return</span> <span class="pl-s1">w</span> <span class="pl-c1">*</span> <span class="pl-s1">x</span> <span class="pl-c1">+</span> <span class="pl-s1">b</span>

<span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">loss_fn</span>(<span class="pl-s1">theta</span>: <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>,
            <span class="pl-s1">x</span>: <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>,
            <span class="pl-s1">y</span>: <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>) <span class="pl-c1">-&gt;</span> <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>:
    <span class="pl-s">'''</span>
<span class="pl-s">    Calculates the MSE loss.</span>
<span class="pl-s">    '''</span>
    <span class="pl-s1">prediction</span> <span class="pl-c1">=</span> <span class="pl-en">model</span>(<span class="pl-s1">theta</span>, <span class="pl-s1">x</span>)            <span class="pl-c"># y_hat</span>
    <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-en">mean</span>((<span class="pl-s1">prediction</span> <span class="pl-c1">-</span> <span class="pl-s1">y</span>) <span class="pl-c1">**</span> <span class="pl-c1">2</span>)  <span class="pl-c"># J</span>

<span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">update</span>(<span class="pl-s1">theta</span>: <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>,
          <span class="pl-s1">x</span>: <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>,
          <span class="pl-s1">y</span>: <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>,
          <span class="pl-s1">lr</span>: <span class="pl-s1">float</span> <span class="pl-c1">=</span> <span class="pl-c1">0.1</span>) <span class="pl-c1">-&gt;</span> <span class="pl-s1">jaxlib</span>.<span class="pl-s1">xla_extension</span>.<span class="pl-v">DeviceArray</span>:
    <span class="pl-s">'''</span>
<span class="pl-s">    updates the parameters based on gradient descent</span>
<span class="pl-s">    '''</span>
    <span class="pl-k">return</span> <span class="pl-s1">theta</span> <span class="pl-c1">-</span> <span class="pl-s1">lr</span> <span class="pl-c1">*</span> <span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-s1">loss_fn</span>)(<span class="pl-s1">theta</span>, <span class="pl-s1">x</span>, <span class="pl-s1">y</span>)

<span class="pl-k">for</span> <span class="pl-s1">_</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">1000</span>):
    <span class="pl-s1">theta</span> <span class="pl-c1">=</span> <span class="pl-en">update</span>(<span class="pl-s1">theta</span>, <span class="pl-s1">x_</span>, <span class="pl-s1">y_</span>)</pre></div>
<p dir="auto">This failed silently.</p>
<p dir="auto">Then I followed the docs and enables <code class="notranslate">nan</code>-parsing by-</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="!JAX_DEBUG_NANS=True"><pre lang="ipython" class="notranslate"><code class="notranslate">!JAX_DEBUG_NANS=True
</code></pre></div>
<p dir="auto">And then I did-</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from jax.config import config
config.update(&quot;jax_debug_nans&quot;, True)"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">config</span> <span class="pl-k">import</span> <span class="pl-s1">config</span>
<span class="pl-s1">config</span>.<span class="pl-en">update</span>(<span class="pl-s">"jax_debug_nans"</span>, <span class="pl-c1">True</span>)</pre></div>
<p dir="auto">Which threw me a error.</p>
<p dir="auto">After the error was thrown, I inspected the <code class="notranslate">theta</code> variable, and saw that it had very small values, i.e. the values before they were going <code class="notranslate">nan</code>.</p>
<p dir="auto">It was this:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; theta
DeviceArray([-1.6214236e+37, -7.4782789e+35], dtype=float32)"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">theta</span>
<span class="pl-v">DeviceArray</span>([<span class="pl-c1">-</span><span class="pl-c1">1.6214236e+37</span>, <span class="pl-c1">-</span><span class="pl-c1">7.4782789e+35</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>)</pre></div>
<p dir="auto">What is going wrong? Am I missing something?</p>
<p dir="auto">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).</p>
<p dir="auto">Full error message:</p>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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
--&gt; 121         for da in tree_leaves(output)
    122         if hasattr(da, &quot;device_buffer&quot;)

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py in check_special(name, bufs)
    409     for buf in bufs:
--&gt; 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())):
--&gt; 416       raise FloatingPointError(f&quot;invalid value (nan) encountered in {name}&quot;)
    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)
--&gt; 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:
--&gt; 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())):
--&gt; 416       raise FloatingPointError(f&quot;invalid value (nan) encountered in {name}&quot;)
    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)
--&gt; 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:
--&gt; 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())):
--&gt; 416       raise FloatingPointError(f&quot;invalid value (nan) encountered in {name}&quot;)
    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)
--&gt; 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:
--&gt; 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())):
--&gt; 416       raise FloatingPointError(f&quot;invalid value (nan) encountered in {name}&quot;)
    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)
--&gt; 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:
--&gt; 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())):
--&gt; 416       raise FloatingPointError(f&quot;invalid value (nan) encountered in {name}&quot;)
    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,
--&gt; 193                      &quot;__main__&quot;, mod_spec)
    194 

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

/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py in &lt;module&gt;
     15     from ipykernel import kernelapp as app
---&gt; 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)
--&gt; 845         app.start()
    846 

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelapp.py in start(***failed resolving arguments***)
    666             try:
--&gt; 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)
--&gt; 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:
--&gt; 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:
-&gt; 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:
---&gt; 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:
--&gt; 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
--&gt; 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):
--&gt; 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):
--&gt; 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
--&gt; 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
--&gt; 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(
-&gt; 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:
-&gt; 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:
---&gt; 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,
-&gt; 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)
-&gt; 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:
-&gt; 3441                     exec(code_obj, self.user_global_ns, self.user_ns)
   3442             finally:

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

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

/tmp/ipykernel_43/3487162561.py in loss_fn(***failed resolving arguments***)
      7     '''
----&gt; 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
----&gt; 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
-&gt; 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)
--&gt; 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 &lt;module&gt;
      1 for _ in range(1000):
----&gt; 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(&quot;Invalid nan value encountered in the output of a C++-jit &quot;
    128           &quot;function. Calling the de-optimized version.&quot;)
--&gt; 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     '''
----&gt; 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())):
--&gt; 416       raise FloatingPointError(f&quot;invalid value (nan) encountered in {name}&quot;)
    417     if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):
    418       raise FloatingPointError(f&quot;invalid value (inf) encountered in {name}&quot;)

FloatingPointError: invalid value (nan) encountered in reduce_sum"><pre class="notranslate">Invalid nan value encountered <span class="pl-k">in</span> the output of a C++-jit function. Calling the de-optimized version.
Invalid value encountered <span class="pl-k">in</span> the output of a jit function. Calling the de-optimized version.
Invalid value encountered <span class="pl-k">in</span> the output of a jit function. Calling the de-optimized version.
Invalid value encountered <span class="pl-k">in</span> the output of a jit function. Calling the de-optimized version.
Invalid value encountered <span class="pl-k">in</span> 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 <span class="pl-k">in</span> _nan_check_posthook(fun, args, kwargs, output)
    120         da.device_buffer
--<span class="pl-k">&gt;</span> 121         <span class="pl-k">for</span> <span class="pl-smi">da</span> <span class="pl-k">in</span> tree_leaves(output)
    122         <span class="pl-k">if</span> hasattr(da, <span class="pl-s"><span class="pl-pds">"</span>device_buffer<span class="pl-pds">"</span></span>)

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py <span class="pl-k">in</span> check_special(name, bufs)
    409     <span class="pl-k">for</span> <span class="pl-smi">buf</span> <span class="pl-k">in</span> bufs:
--<span class="pl-k">&gt;</span> 410       _check_special(name, <span class="pl-en">buf.xla_shape</span>(), buf)
    411 

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py <span class="pl-k">in</span> _check_special(name, xla_shape, buf)
    415     <span class="pl-k">if</span> config.jax_debug_nans and <span class="pl-en">np.any(np.isnan(buf.to_py</span>())):
--<span class="pl-k">&gt;</span> 416       raise FloatingPointError(f<span class="pl-s"><span class="pl-pds">"</span>invalid value (nan) encountered in {name}<span class="pl-pds">"</span></span>)
    417     <span class="pl-k">if</span> config.jax_debug_infs and <span class="pl-en">np.any(np.isinf(buf.to_py</span>())):

FloatingPointError: invalid value (nan) encountered <span class="pl-k">in</span> 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 <span class="pl-k">in</span> _execute_compiled(compiled, avals, handlers, kept_var_idx, <span class="pl-k">*</span>args)
    912   out_bufs = compiled.execute(input_bufs)
--<span class="pl-k">&gt;</span> 913   check_special(xla_call_p.name, out_bufs)
    914   <span class="pl-k">return</span> [handler(<span class="pl-k">*</span>bs) <span class="pl-k">for</span> <span class="pl-smi">handler, bs</span> <span class="pl-k">in</span> zip(handlers, _partition_outputs(avals, out_bufs))]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py <span class="pl-k">in</span> check_special(name, bufs)
    409     <span class="pl-k">for</span> <span class="pl-smi">buf</span> <span class="pl-k">in</span> bufs:
--<span class="pl-k">&gt;</span> 410       _check_special(name, <span class="pl-en">buf.xla_shape</span>(), buf)
    411 

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py <span class="pl-k">in</span> _check_special(name, xla_shape, buf)
    415     <span class="pl-k">if</span> config.jax_debug_nans and <span class="pl-en">np.any(np.isnan(buf.to_py</span>())):
--<span class="pl-k">&gt;</span> 416       raise FloatingPointError(f<span class="pl-s"><span class="pl-pds">"</span>invalid value (nan) encountered in {name}<span class="pl-pds">"</span></span>)
    417     <span class="pl-k">if</span> config.jax_debug_infs and <span class="pl-en">np.any(np.isinf(buf.to_py</span>())):

FloatingPointError: invalid value (nan) encountered <span class="pl-k">in</span> 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 <span class="pl-k">in</span> _execute_compiled(compiled, avals, handlers, kept_var_idx, <span class="pl-k">*</span>args)
    912   out_bufs = compiled.execute(input_bufs)
--<span class="pl-k">&gt;</span> 913   check_special(xla_call_p.name, out_bufs)
    914   <span class="pl-k">return</span> [handler(<span class="pl-k">*</span>bs) <span class="pl-k">for</span> <span class="pl-smi">handler, bs</span> <span class="pl-k">in</span> zip(handlers, _partition_outputs(avals, out_bufs))]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py <span class="pl-k">in</span> check_special(name, bufs)
    409     <span class="pl-k">for</span> <span class="pl-smi">buf</span> <span class="pl-k">in</span> bufs:
--<span class="pl-k">&gt;</span> 410       _check_special(name, <span class="pl-en">buf.xla_shape</span>(), buf)
    411 

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py <span class="pl-k">in</span> _check_special(name, xla_shape, buf)
    415     <span class="pl-k">if</span> config.jax_debug_nans and <span class="pl-en">np.any(np.isnan(buf.to_py</span>())):
--<span class="pl-k">&gt;</span> 416       raise FloatingPointError(f<span class="pl-s"><span class="pl-pds">"</span>invalid value (nan) encountered in {name}<span class="pl-pds">"</span></span>)
    417     <span class="pl-k">if</span> config.jax_debug_infs and <span class="pl-en">np.any(np.isinf(buf.to_py</span>())):

FloatingPointError: invalid value (nan) encountered <span class="pl-k">in</span> 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 <span class="pl-k">in</span> _execute_compiled(compiled, avals, handlers, kept_var_idx, <span class="pl-k">*</span>args)
    912   out_bufs = compiled.execute(input_bufs)
--<span class="pl-k">&gt;</span> 913   check_special(xla_call_p.name, out_bufs)
    914   <span class="pl-k">return</span> [handler(<span class="pl-k">*</span>bs) <span class="pl-k">for</span> <span class="pl-smi">handler, bs</span> <span class="pl-k">in</span> zip(handlers, _partition_outputs(avals, out_bufs))]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py <span class="pl-k">in</span> check_special(name, bufs)
    409     <span class="pl-k">for</span> <span class="pl-smi">buf</span> <span class="pl-k">in</span> bufs:
--<span class="pl-k">&gt;</span> 410       _check_special(name, <span class="pl-en">buf.xla_shape</span>(), buf)
    411 

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py <span class="pl-k">in</span> _check_special(name, xla_shape, buf)
    415     <span class="pl-k">if</span> config.jax_debug_nans and <span class="pl-en">np.any(np.isnan(buf.to_py</span>())):
--<span class="pl-k">&gt;</span> 416       raise FloatingPointError(f<span class="pl-s"><span class="pl-pds">"</span>invalid value (nan) encountered in {name}<span class="pl-pds">"</span></span>)
    417     <span class="pl-k">if</span> config.jax_debug_infs and <span class="pl-en">np.any(np.isinf(buf.to_py</span>())):

FloatingPointError: invalid value (nan) encountered <span class="pl-k">in</span> 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 <span class="pl-k">in</span> _execute_compiled(compiled, avals, handlers, kept_var_idx, <span class="pl-k">*</span>args)
    912   out_bufs = compiled.execute(input_bufs)
--<span class="pl-k">&gt;</span> 913   check_special(xla_call_p.name, out_bufs)
    914   <span class="pl-k">return</span> [handler(<span class="pl-k">*</span>bs) <span class="pl-k">for</span> <span class="pl-smi">handler, bs</span> <span class="pl-k">in</span> zip(handlers, _partition_outputs(avals, out_bufs))]

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py <span class="pl-k">in</span> check_special(name, bufs)
    409     <span class="pl-k">for</span> <span class="pl-smi">buf</span> <span class="pl-k">in</span> bufs:
--<span class="pl-k">&gt;</span> 410       _check_special(name, <span class="pl-en">buf.xla_shape</span>(), buf)
    411 

/opt/conda/lib/python3.7/site-packages/jax/interpreters/xla.py <span class="pl-k">in</span> _check_special(name, xla_shape, buf)
    415     <span class="pl-k">if</span> config.jax_debug_nans and <span class="pl-en">np.any(np.isnan(buf.to_py</span>())):
--<span class="pl-k">&gt;</span> 416       raise FloatingPointError(f<span class="pl-s"><span class="pl-pds">"</span>invalid value (nan) encountered in {name}<span class="pl-pds">"</span></span>)
    417     <span class="pl-k">if</span> config.jax_debug_infs and <span class="pl-en">np.any(np.isinf(buf.to_py</span>())):

FloatingPointError: invalid value (nan) encountered <span class="pl-k">in</span> xla_call

During handling of the above exception, another exception occurred:

JaxStackTraceBeforeTransformation         Traceback (most recent call last)
/opt/conda/lib/python3.7/runpy.py <span class="pl-k">in</span> _run_module_as_main(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
    192     <span class="pl-k">return</span> _run_code(code, main_globals, None,
--<span class="pl-k">&gt;</span> 193                      <span class="pl-s"><span class="pl-pds">"</span>__main__<span class="pl-pds">"</span></span>, mod_spec)
    194 

/opt/conda/lib/python3.7/runpy.py <span class="pl-k">in</span> _run_code(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
     84                        __spec__ = mod_spec)
---<span class="pl-k">&gt;</span> 85     exec(code, run_globals)
     86     <span class="pl-k">return</span> run_globals

/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py <span class="pl-k">in</span> <span class="pl-k">&lt;</span>module<span class="pl-k">&gt;</span>
     15     from ipykernel import kernelapp as app
---<span class="pl-k">&gt;</span> 16     <span class="pl-en">app.launch_new_instance</span>()

/opt/conda/lib/python3.7/site-packages/traitlets/config/application.py <span class="pl-k">in</span> launch_instance(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
    844         app.initialize(argv)
--<span class="pl-k">&gt;</span> 845         <span class="pl-en">app.start</span>()
    846 

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelapp.py <span class="pl-k">in</span> start(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
    666             try:
--<span class="pl-k">&gt;</span> 667                 <span class="pl-en">self.io_loop.start</span>()
    668             except KeyboardInterrupt:

/opt/conda/lib/python3.7/site-packages/tornado/platform/asyncio.py <span class="pl-k">in</span> start(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
    198             asyncio.set_event_loop(self.asyncio_loop)
--<span class="pl-k">&gt;</span> 199             <span class="pl-en">self.asyncio_loop.run_forever</span>()
    200         finally:

/opt/conda/lib/python3.7/asyncio/base_events.py <span class="pl-k">in</span> run_forever(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
    540             <span class="pl-k">while</span> True:
--<span class="pl-k">&gt;</span> 541                 <span class="pl-en">self._run_once</span>()
    542                 <span class="pl-k">if</span> self._stopping:

/opt/conda/lib/python3.7/asyncio/base_events.py <span class="pl-k">in</span> _run_once(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
   1785             else:
-<span class="pl-k">&gt;</span> 1786                 <span class="pl-en">handle._run</span>()
   1787         handle = None  <span class="pl-c"><span class="pl-c">#</span> Needed to break cycles when an exception occurs.</span>

/opt/conda/lib/python3.7/asyncio/events.py <span class="pl-k">in</span> _run(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
     87         try:
---<span class="pl-k">&gt;</span> 88             self._context.run(self._callback, <span class="pl-k">*</span>self._args)
     89         except Exception as exc:

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py <span class="pl-k">in</span> dispatch_queue(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
    456             try:
--<span class="pl-k">&gt;</span> 457                 await <span class="pl-en">self.process_one</span>()
    458             except Exception:

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py <span class="pl-k">in</span> process_one(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
    445                 <span class="pl-k">return</span> None
--<span class="pl-k">&gt;</span> 446         await dispatch(<span class="pl-k">*</span>args)
    447 

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py <span class="pl-k">in</span> dispatch_shell(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
    352                 <span class="pl-k">if</span> inspect.isawaitable(result):
--<span class="pl-k">&gt;</span> 353                     await result
    354             except Exception:

/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py <span class="pl-k">in</span> execute_request(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
    647         <span class="pl-k">if</span> inspect.isawaitable(reply_content):
--<span class="pl-k">&gt;</span> 648             reply_content = await reply_content
    649 

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

/opt/conda/lib/python3.7/site-packages/ipykernel/zmqshell.py <span class="pl-k">in</span> run_cell(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
    531         self._last_traceback = None
--<span class="pl-k">&gt;</span> 532         <span class="pl-k">return</span> super(ZMQInteractiveShell, self).run_cell(<span class="pl-k">*</span>args, <span class="pl-k">**</span>kwargs)
    533 

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

/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py <span class="pl-k">in</span> _run_cell(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
   2943         try:
-<span class="pl-k">&gt;</span> 2944             <span class="pl-k">return</span> runner(coro)
   2945         except BaseException as e:

/opt/conda/lib/python3.7/site-packages/IPython/core/async_helpers.py <span class="pl-k">in</span> _pseudo_sync_runner(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
     67     try:
---<span class="pl-k">&gt;</span> 68         coro.send(None)
     69     except StopIteration as exc:

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

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

/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py <span class="pl-k">in</span> run_code(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
   3440                 else:
-<span class="pl-k">&gt;</span> 3441                     exec(code_obj, self.user_global_ns, self.user_ns)
   3442             finally:

/tmp/ipykernel_43/2138327653.py <span class="pl-k">in</span> <span class="pl-k">&lt;</span>module<span class="pl-k">&gt;</span>
      1 <span class="pl-k">for</span> <span class="pl-smi">_</span> <span class="pl-k">in</span> range(1000):
----<span class="pl-k">&gt;</span> 2     theta = update(theta, x_train, y_train)
      3 

/tmp/ipykernel_43/1444374627.py <span class="pl-k">in</span> update(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)
      8     <span class="pl-s"><span class="pl-pds">'</span><span class="pl-pds">'</span><span class="pl-pds">'</span></span>
<span class="pl-s">----&gt; 9     return theta - lr * jax.grad(loss_fn)(theta, x, y)</span>
<span class="pl-s"></span>
<span class="pl-s">/tmp/ipykernel_43/3487162561.py in loss_fn(***failed resolving arguments***)</span>
<span class="pl-s">      7     <span class="pl-pds">'</span><span class="pl-pds">'</span><span class="pl-pds">'</span></span>
----<span class="pl-k">&gt;</span> 8     prediction = model(theta, x)            <span class="pl-c"><span class="pl-c">#</span> y_hat</span>
      9     <span class="pl-k">return</span> jnp.mean<span class="pl-s"><span class="pl-pds">((</span>prediction <span class="pl-k">-</span> y) <span class="pl-k">**</span> <span class="pl-c1">2</span>)  # J</span>
<span class="pl-s"></span>
<span class="pl-s"><span class="pl-k">/</span>tmp<span class="pl-k">/</span>ipykernel_<span class="pl-c1">43</span><span class="pl-k">/</span><span class="pl-c1">3906416787</span>.py in model(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)</span>
<span class="pl-s">      <span class="pl-c1">3</span>     w, b = theta</span>
<span class="pl-s"><span class="pl-k">----&gt;</span> <span class="pl-c1">4</span>     return w <span class="pl-k">*</span> x <span class="pl-k">+</span> b</span>
<span class="pl-s"></span>
<span class="pl-s"><span class="pl-k">/</span>opt<span class="pl-k">/</span>conda<span class="pl-k">/</span>lib<span class="pl-k">/</span>python<span class="pl-c1">3</span>.<span class="pl-c1">7</span><span class="pl-k">/</span>site<span class="pl-k">-</span>packages<span class="pl-k">/</span>jax<span class="pl-k">/</span>_src<span class="pl-k">/</span>numpy<span class="pl-k">/</span>lax_numpy.py in deferring_binary_op(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)</span>
<span class="pl-s">   <span class="pl-c1">5923</span>       return NotImplemented</span>
<span class="pl-s"><span class="pl-k">-&gt;</span> <span class="pl-c1">5924</span>     return binary_op(self<span class="pl-k">,</span> other)</span>
<span class="pl-s">   <span class="pl-c1">5925</span>   return deferring_binary_op</span>
<span class="pl-s"></span>
<span class="pl-s"><span class="pl-k">/</span>opt<span class="pl-k">/</span>conda<span class="pl-k">/</span>lib<span class="pl-k">/</span>python<span class="pl-c1">3</span>.<span class="pl-c1">7</span><span class="pl-k">/</span>site<span class="pl-k">-</span>packages<span class="pl-k">/</span>jax<span class="pl-k">/</span>_src<span class="pl-k">/</span>numpy<span class="pl-k">/</span>lax_numpy.py in fn(<span class="pl-k">***</span>failed resolving arguments<span class="pl-k">***</span>)</span>
<span class="pl-s">    <span class="pl-c1">435</span>     x<span class="pl-c1">1</span>, x<span class="pl-c1">2</span> = _promote_args(numpy_fn.__name__, x<span class="pl-c1">1</span>, x<span class="pl-c1">2</span>)</span>
<span class="pl-s"><span class="pl-k">--&gt;</span> <span class="pl-c1">436</span>     return lax_fn(x<span class="pl-c1">1</span><span class="pl-k">,</span> x<span class="pl-c1">2</span>) if x<span class="pl-c1">1</span>.dtype <span class="pl-k">!=</span> bool_ else bool_lax_fn(x<span class="pl-c1">1</span><span class="pl-k">,</span> x<span class="pl-c1">2</span>)</span>
<span class="pl-s">    <span class="pl-c1">437</span>   fn = jit(fn, inline=True)</span>
<span class="pl-s"></span>
<span class="pl-s">JaxStackTraceBeforeTransformation: FloatingPointError: invalid value (nan) encountered in reduce_sum</span>
<span class="pl-s"></span>
<span class="pl-s">The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.</span>
<span class="pl-s"></span>
<span class="pl-s"><span class="pl-k">--------------------</span></span>
<span class="pl-s"></span>
<span class="pl-s">The above exception was the direct cause of the following exception:</span>
<span class="pl-s"></span>
<span class="pl-s">FloatingPointError                        Traceback (most recent call last)</span>
<span class="pl-s"><span class="pl-k">/</span>tmp<span class="pl-k">/</span>ipykernel_<span class="pl-c1">43</span><span class="pl-k">/</span><span class="pl-c1">2138327653</span>.py in <span class="pl-k">&lt;</span>module<span class="pl-k">&gt;</span></span>
<span class="pl-s">      <span class="pl-c1">1</span> for _ in range(<span class="pl-c1">1000</span>):</span>
<span class="pl-s"><span class="pl-k">----&gt;</span> <span class="pl-c1">2</span>     theta <span class="pl-k">=</span> update(theta<span class="pl-k">,</span> x_train<span class="pl-k">,</span> y_train)</span>
<span class="pl-s">      <span class="pl-c1">3</span> </span>
<span class="pl-s">      <span class="pl-c1">4</span> plt.scatter(x_train, y_train)</span>
<span class="pl-s">      <span class="pl-c1">5</span> plt.plot(x_train, model(theta, x_train), 'r')</span>
<span class="pl-s"></span>
<span class="pl-s"><span class="pl-k">/</span>opt<span class="pl-k">/</span>conda<span class="pl-k">/</span>lib<span class="pl-k">/</span>python<span class="pl-c1">3</span>.<span class="pl-c1">7</span><span class="pl-k">/</span>site<span class="pl-k">-</span>packages<span class="pl-k">/</span>jax<span class="pl-k">/</span>_src<span class="pl-k">/</span>api.py in _nan_check_posthook(fun<span class="pl-k">,</span> args<span class="pl-k">,</span> kwargs<span class="pl-k">,</span> output)</span>
<span class="pl-s">    <span class="pl-c1">127</span>     print("Invalid nan value encountered in the output of a C++-jit "</span>
<span class="pl-s">    <span class="pl-c1">128</span>           "function. Calling the de-optimized version.")</span>
<span class="pl-s"><span class="pl-k">--&gt;</span> <span class="pl-c1">129</span>     fun._cache_miss(<span class="pl-k">*</span>args<span class="pl-k">,</span> <span class="pl-k">**</span>kwargs)[<span class="pl-c1">0</span>]  # probably won't return</span>
<span class="pl-s">    <span class="pl-c1">130</span> </span>
<span class="pl-s">    <span class="pl-c1">131</span> def _update_debug_special_global(_):</span>
<span class="pl-s"></span>
<span class="pl-s">    [... skipping hidden <span class="pl-c1">8</span> frame]</span>
<span class="pl-s"></span>
<span class="pl-s"><span class="pl-k">/</span>tmp<span class="pl-k">/</span>ipykernel_<span class="pl-c1">43</span><span class="pl-k">/</span><span class="pl-c1">1444374627</span>.py in update(theta<span class="pl-k">,</span> x<span class="pl-k">,</span> y<span class="pl-k">,</span> lr)</span>
<span class="pl-s">      <span class="pl-c1">7</span>     updates the parameters based on gradient descent</span>
<span class="pl-s">      <span class="pl-c1">8</span>     '''</span>
<span class="pl-s"><span class="pl-k">----&gt;</span> <span class="pl-c1">9</span>     return theta <span class="pl-k">-</span> lr <span class="pl-k">*</span> jax.grad(loss_fn)(theta<span class="pl-k">,</span> x<span class="pl-k">,</span> y)</span>
<span class="pl-s"></span>
<span class="pl-s">    [... skipping hidden <span class="pl-c1">39</span> frame]</span>
<span class="pl-s"></span>
<span class="pl-s"><span class="pl-k">/</span>opt<span class="pl-k">/</span>conda<span class="pl-k">/</span>lib<span class="pl-k">/</span>python<span class="pl-c1">3</span>.<span class="pl-c1">7</span><span class="pl-k">/</span>site<span class="pl-k">-</span>packages<span class="pl-k">/</span>jax<span class="pl-k">/</span>interpreters<span class="pl-k">/</span>xla.py in _check_special(name<span class="pl-k">,</span> xla_shape<span class="pl-k">,</span> buf)</span>
<span class="pl-s">    <span class="pl-c1">414</span>   if dtypes.issubdtype(xla_shape.element_type(), np.inexact):</span>
<span class="pl-s">    <span class="pl-c1">415</span>     if config.jax_debug_nans and np.any(np.isnan(buf.to_py(<span class="pl-pds">))</span></span>):
--<span class="pl-k">&gt;</span> 416       raise FloatingPointError(f<span class="pl-s"><span class="pl-pds">"</span>invalid value (nan) encountered in {name}<span class="pl-pds">"</span></span>)
    417     <span class="pl-k">if</span> config.jax_debug_infs and <span class="pl-en">np.any(np.isinf(buf.to_py</span>())):
    418       raise FloatingPointError(f<span class="pl-s"><span class="pl-pds">"</span>invalid value (inf) encountered in {name}<span class="pl-pds">"</span></span>)

FloatingPointError: invalid value (nan) encountered <span class="pl-k">in</span> reduce_sum</pre></div>