lars/FYS-STK4155
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
Morten Hjorth-Jensen
2023-08-31 06:19:44 +02:00
parent e9bac09427
commit d05f791ff6
98 changed files with 3861 additions and 2496 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/LectureNotes/_build/html/chapter12.html
+147 -8
View File
@@ -242,6 +242,33 @@ const thebe_selector_output = ".output, .cell_output"
</a>
</li>
</ul>
<p aria-level="2" class="caption" role="heading">
<span class="caption-text">
Weekly material, notes and exercises
</span>
</p>
<ul class="nav bd-sidenav">
<li class="toctree-l1">
<a class="reference internal" href="exercisesweek34.html">
Exercises week 34
</a>
</li>
<li class="toctree-l1">
<a class="reference internal" href="week34.html">
Week 34: Introduction to the course, Logistics and Practicalities
</a>
</li>
<li class="toctree-l1">
<a class="reference internal" href="exercisesweek35.html">
Exercises week 35
</a>
</li>
<li class="toctree-l1">
<a class="reference internal" href="week35.html">
Week 35: From Ordinary Linear Regression to Ridge and Lasso Regression
</a>
</li>
</ul>
</div>
</nav> <!-- To handle the deprecated key -->
@@ -1129,8 +1156,12 @@ labels = (n_inputs) = (1797,)
</div>
</div>
<div class="cell_output docutils container">
<div class="output stderr highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>/Users/mhjensen/miniforge3/envs/myenv/lib/python3.9/site-packages/jax/_src/lib/__init__.py:32: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems.
warnings.warn(&quot;JAX on Mac ARM machines is experimental and minimally tested. &quot;
</pre></div>
</div>
<div class="output traceback highlight-ipythontb notranslate"><div class="highlight"><pre><span></span><span class="gt">---------------------------------------------------------------------------</span>
<span class="ne">ModuleNotFoundError</span><span class="g g-Whitespace"> </span>Traceback (most recent call last)
<span class="ne">AttributeError</span><span class="g g-Whitespace"> </span>Traceback (most recent call last)
<span class="nn">Input In [4],</span> in <span class="ni">&lt;cell line: 1&gt;</span><span class="nt">()</span>
<span class="ne">----&gt; </span><span class="mi">1</span> <span class="kn">from</span> <span class="nn">tensorflow.keras</span> <span class="kn">import</span> <span class="n">datasets</span><span class="p">,</span> <span class="n">layers</span><span class="p">,</span> <span class="n">models</span>
<span class="g g-Whitespace"> </span><span class="mi">2</span> <span class="kn">from</span> <span class="nn">tensorflow.keras.layers</span> <span class="kn">import</span> <span class="n">Input</span>
@@ -1216,14 +1247,122 @@ labels = (n_inputs) = (1797,)
<span class="g g-Whitespace"> </span><span class="mi">30</span> <span class="kn">from</span> <span class="nn">tensorflow.lite.python</span> <span class="kn">import</span> <span class="n">wrap_toco</span>
<span class="g g-Whitespace"> </span><span class="mi">31</span> <span class="kn">from</span> <span class="nn">tensorflow.lite.python.convert_phase</span> <span class="kn">import</span> <span class="n">Component</span>
<span class="n">File</span> <span class="o">~/</span><span class="n">miniforge3</span><span class="o">/</span><span class="n">envs</span><span class="o">/</span><span class="n">myenv</span><span class="o">/</span><span class="n">lib</span><span class="o">/</span><span class="n">python3</span><span class="mf">.9</span><span class="o">/</span><span class="n">site</span><span class="o">-</span><span class="n">packages</span><span class="o">/</span><span class="n">tensorflow</span><span class="o">/</span><span class="n">lite</span><span class="o">/</span><span class="n">python</span><span class="o">/</span><span class="n">util</span><span class="o">.</span><span class="n">py</span><span class="p">:</span><span class="mi">26</span><span class="p">,</span> <span class="ow">in</span> <span class="o">&lt;</span><span class="n">module</span><span class="o">&gt;</span>
<span class="g g-Whitespace"> </span><span class="mi">23</span> <span class="kn">import</span> <span class="nn">six</span>
<span class="g g-Whitespace"> </span><span class="mi">24</span> <span class="kn">from</span> <span class="nn">six.moves</span> <span class="kn">import</span> <span class="nb">range</span>
<span class="ne">---&gt; </span><span class="mi">26</span> <span class="kn">import</span> <span class="nn">flatbuffers</span>
<span class="g g-Whitespace"> </span><span class="mi">27</span> <span class="kn">from</span> <span class="nn">tensorflow.core.protobuf</span> <span class="kn">import</span> <span class="n">config_pb2</span> <span class="k">as</span> <span class="n">_config_pb2</span>
<span class="g g-Whitespace"> </span><span class="mi">28</span> <span class="kn">from</span> <span class="nn">tensorflow.core.protobuf</span> <span class="kn">import</span> <span class="n">graph_debug_info_pb2</span>
<span class="n">File</span> <span class="o">~/</span><span class="n">miniforge3</span><span class="o">/</span><span class="n">envs</span><span class="o">/</span><span class="n">myenv</span><span class="o">/</span><span class="n">lib</span><span class="o">/</span><span class="n">python3</span><span class="mf">.9</span><span class="o">/</span><span class="n">site</span><span class="o">-</span><span class="n">packages</span><span class="o">/</span><span class="n">tensorflow</span><span class="o">/</span><span class="n">lite</span><span class="o">/</span><span class="n">python</span><span class="o">/</span><span class="n">util</span><span class="o">.</span><span class="n">py</span><span class="p">:</span><span class="mi">51</span><span class="p">,</span> <span class="ow">in</span> <span class="o">&lt;</span><span class="n">module</span><span class="o">&gt;</span>
<span class="g g-Whitespace"> </span><span class="mi">47</span> <span class="c1"># Jax functions used by TFLite</span>
<span class="g g-Whitespace"> </span><span class="mi">48</span> <span class="c1"># pylint: disable=g-import-not-at-top</span>
<span class="g g-Whitespace"> </span><span class="mi">49</span> <span class="c1"># pylint: disable=unused-import</span>
<span class="g g-Whitespace"> </span><span class="mi">50</span> <span class="k">try</span><span class="p">:</span>
<span class="ne">---&gt; </span><span class="mi">51</span> <span class="kn">from</span> <span class="nn">jax</span> <span class="kn">import</span> <span class="n">xla_computation</span> <span class="k">as</span> <span class="n">_xla_computation</span>
<span class="g g-Whitespace"> </span><span class="mi">52</span> <span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span>
<span class="g g-Whitespace"> </span><span class="mi">53</span> <span class="n">_xla_computation</span> <span class="o">=</span> <span class="kc">None</span>
<span class="ne">ModuleNotFoundError</span>: No module named &#39;flatbuffers&#39;
<span class="n">File</span> <span class="o">~/</span><span class="n">miniforge3</span><span class="o">/</span><span class="n">envs</span><span class="o">/</span><span class="n">myenv</span><span class="o">/</span><span class="n">lib</span><span class="o">/</span><span class="n">python3</span><span class="mf">.9</span><span class="o">/</span><span class="n">site</span><span class="o">-</span><span class="n">packages</span><span class="o">/</span><span class="n">jax</span><span class="o">/</span><span class="fm">__init__</span><span class="o">.</span><span class="n">py</span><span class="p">:</span><span class="mi">116</span><span class="p">,</span> <span class="ow">in</span> <span class="o">&lt;</span><span class="n">module</span><span class="o">&gt;</span>
<span class="g g-Whitespace"> </span><span class="mi">40</span> <span class="kn">from</span> <span class="nn">._src.config</span> <span class="kn">import</span> <span class="p">(</span>
<span class="g g-Whitespace"> </span><span class="mi">41</span> <span class="n">config</span> <span class="k">as</span> <span class="n">config</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">42</span> <span class="n">enable_checks</span> <span class="k">as</span> <span class="n">enable_checks</span><span class="p">,</span>
<span class="p">(</span><span class="o">...</span><span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">51</span> <span class="n">numpy_rank_promotion</span> <span class="k">as</span> <span class="n">numpy_rank_promotion</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">52</span> <span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">53</span> <span class="kn">from</span> <span class="nn">._src.api</span> <span class="kn">import</span> <span class="p">(</span>
<span class="g g-Whitespace"> </span><span class="mi">54</span> <span class="n">ad</span><span class="p">,</span> <span class="c1"># TODO(phawkins): update users to avoid this.</span>
<span class="g g-Whitespace"> </span><span class="mi">55</span> <span class="n">checkpoint</span> <span class="k">as</span> <span class="n">checkpoint</span><span class="p">,</span>
<span class="p">(</span><span class="o">...</span><span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">114</span> <span class="n">xla_computation</span> <span class="k">as</span> <span class="n">xla_computation</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">115</span> <span class="p">)</span>
<span class="ne">--&gt; </span><span class="mi">116</span> <span class="kn">from</span> <span class="nn">.experimental.maps</span> <span class="kn">import</span> <span class="n">soft_pmap</span> <span class="k">as</span> <span class="n">soft_pmap</span>
<span class="g g-Whitespace"> </span><span class="mi">117</span> <span class="kn">from</span> <span class="nn">.version</span> <span class="kn">import</span> <span class="n">__version__</span> <span class="k">as</span> <span class="n">__version__</span>
<span class="g g-Whitespace"> </span><span class="mi">119</span> <span class="c1"># These submodules are separate because they are in an import cycle with</span>
<span class="g g-Whitespace"> </span><span class="mi">120</span> <span class="c1"># jax and rely on the names imported above.</span>
<span class="n">File</span> <span class="o">~/</span><span class="n">miniforge3</span><span class="o">/</span><span class="n">envs</span><span class="o">/</span><span class="n">myenv</span><span class="o">/</span><span class="n">lib</span><span class="o">/</span><span class="n">python3</span><span class="mf">.9</span><span class="o">/</span><span class="n">site</span><span class="o">-</span><span class="n">packages</span><span class="o">/</span><span class="n">jax</span><span class="o">/</span><span class="n">experimental</span><span class="o">/</span><span class="n">maps</span><span class="o">.</span><span class="n">py</span><span class="p">:</span><span class="mi">26</span><span class="p">,</span> <span class="ow">in</span> <span class="o">&lt;</span><span class="n">module</span><span class="o">&gt;</span>
<span class="g g-Whitespace"> </span><span class="mi">23</span> <span class="kn">from</span> <span class="nn">functools</span> <span class="kn">import</span> <span class="n">wraps</span><span class="p">,</span> <span class="n">partial</span><span class="p">,</span> <span class="n">partialmethod</span>
<span class="g g-Whitespace"> </span><span class="mi">24</span> <span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
<span class="ne">---&gt; </span><span class="mi">26</span> <span class="kn">from</span> <span class="nn">..</span> <span class="kn">import</span> <span class="n">numpy</span> <span class="k">as</span> <span class="n">jnp</span>
<span class="g g-Whitespace"> </span><span class="mi">27</span> <span class="kn">from</span> <span class="nn">..</span> <span class="kn">import</span> <span class="n">core</span>
<span class="g g-Whitespace"> </span><span class="mi">28</span> <span class="kn">from</span> <span class="nn">..</span> <span class="kn">import</span> <span class="n">linear_util</span> <span class="k">as</span> <span class="n">lu</span>
<span class="n">File</span> <span class="o">~/</span><span class="n">miniforge3</span><span class="o">/</span><span class="n">envs</span><span class="o">/</span><span class="n">myenv</span><span class="o">/</span><span class="n">lib</span><span class="o">/</span><span class="n">python3</span><span class="mf">.9</span><span class="o">/</span><span class="n">site</span><span class="o">-</span><span class="n">packages</span><span class="o">/</span><span class="n">jax</span><span class="o">/</span><span class="n">numpy</span><span class="o">/</span><span class="fm">__init__</span><span class="o">.</span><span class="n">py</span><span class="p">:</span><span class="mi">19</span><span class="p">,</span> <span class="ow">in</span> <span class="o">&lt;</span><span class="n">module</span><span class="o">&gt;</span>
<span class="g g-Whitespace"> </span><span class="mi">1</span> <span class="c1"># Copyright 2018 Google LLC</span>
<span class="g g-Whitespace"> </span><span class="mi">2</span> <span class="c1">#</span>
<span class="g g-Whitespace"> </span><span class="mi">3</span> <span class="c1"># Licensed under the Apache License, Version 2.0 (the &quot;License&quot;);</span>
<span class="p">(</span><span class="o">...</span><span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">17</span>
<span class="g g-Whitespace"> </span><span class="mi">18</span> <span class="c1"># flake8: noqa: F401</span>
<span class="ne">---&gt; </span><span class="mi">19</span> <span class="kn">from</span> <span class="nn">.</span> <span class="kn">import</span> <span class="n">fft</span> <span class="k">as</span> <span class="n">fft</span>
<span class="g g-Whitespace"> </span><span class="mi">20</span> <span class="kn">from</span> <span class="nn">.</span> <span class="kn">import</span> <span class="n">linalg</span> <span class="k">as</span> <span class="n">linalg</span>
<span class="g g-Whitespace"> </span><span class="mi">22</span> <span class="kn">from</span> <span class="nn">jax.interpreters.xla</span> <span class="kn">import</span> <span class="n">DeviceArray</span> <span class="k">as</span> <span class="n">DeviceArray</span>
<span class="n">File</span> <span class="o">~/</span><span class="n">miniforge3</span><span class="o">/</span><span class="n">envs</span><span class="o">/</span><span class="n">myenv</span><span class="o">/</span><span class="n">lib</span><span class="o">/</span><span class="n">python3</span><span class="mf">.9</span><span class="o">/</span><span class="n">site</span><span class="o">-</span><span class="n">packages</span><span class="o">/</span><span class="n">jax</span><span class="o">/</span><span class="n">numpy</span><span class="o">/</span><span class="n">fft</span><span class="o">.</span><span class="n">py</span><span class="p">:</span><span class="mi">17</span><span class="p">,</span> <span class="ow">in</span> <span class="o">&lt;</span><span class="n">module</span><span class="o">&gt;</span>
<span class="g g-Whitespace"> </span><span class="mi">1</span> <span class="c1"># Copyright 2020 Google LLC</span>
<span class="g g-Whitespace"> </span><span class="mi">2</span> <span class="c1">#</span>
<span class="g g-Whitespace"> </span><span class="mi">3</span> <span class="c1"># Licensed under the Apache License, Version 2.0 (the &quot;License&quot;);</span>
<span class="p">(</span><span class="o">...</span><span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">14</span>
<span class="g g-Whitespace"> </span><span class="mi">15</span> <span class="c1"># flake8: noqa: F401</span>
<span class="ne">---&gt; </span><span class="mi">17</span> <span class="kn">from</span> <span class="nn">jax._src.numpy.fft</span> <span class="kn">import</span> <span class="p">(</span>
<span class="g g-Whitespace"> </span><span class="mi">18</span> <span class="n">ifft</span> <span class="k">as</span> <span class="n">ifft</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">19</span> <span class="n">ifft2</span> <span class="k">as</span> <span class="n">ifft2</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">20</span> <span class="n">ifftn</span> <span class="k">as</span> <span class="n">ifftn</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">21</span> <span class="n">ifftshift</span> <span class="k">as</span> <span class="n">ifftshift</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">22</span> <span class="n">ihfft</span> <span class="k">as</span> <span class="n">ihfft</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">23</span> <span class="n">irfft</span> <span class="k">as</span> <span class="n">irfft</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">24</span> <span class="n">irfft2</span> <span class="k">as</span> <span class="n">irfft2</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">25</span> <span class="n">irfftn</span> <span class="k">as</span> <span class="n">irfftn</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">26</span> <span class="n">fft</span> <span class="k">as</span> <span class="n">fft</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">27</span> <span class="n">fft2</span> <span class="k">as</span> <span class="n">fft2</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">28</span> <span class="n">fftfreq</span> <span class="k">as</span> <span class="n">fftfreq</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">29</span> <span class="n">fftn</span> <span class="k">as</span> <span class="n">fftn</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">30</span> <span class="n">fftshift</span> <span class="k">as</span> <span class="n">fftshift</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">31</span> <span class="n">hfft</span> <span class="k">as</span> <span class="n">hfft</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">32</span> <span class="n">rfft</span> <span class="k">as</span> <span class="n">rfft</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">33</span> <span class="n">rfft2</span> <span class="k">as</span> <span class="n">rfft2</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">34</span> <span class="n">rfftfreq</span> <span class="k">as</span> <span class="n">rfftfreq</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">35</span> <span class="n">rfftn</span> <span class="k">as</span> <span class="n">rfftn</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">36</span> <span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">38</span> <span class="c1"># Module initialization is encapsulated in a function to avoid accidental</span>
<span class="g g-Whitespace"> </span><span class="mi">39</span> <span class="c1"># namespace pollution.</span>
<span class="g g-Whitespace"> </span><span class="mi">40</span> <span class="n">_NOT_IMPLEMENTED</span> <span class="o">=</span> <span class="p">[]</span>
<span class="n">File</span> <span class="o">~/</span><span class="n">miniforge3</span><span class="o">/</span><span class="n">envs</span><span class="o">/</span><span class="n">myenv</span><span class="o">/</span><span class="n">lib</span><span class="o">/</span><span class="n">python3</span><span class="mf">.9</span><span class="o">/</span><span class="n">site</span><span class="o">-</span><span class="n">packages</span><span class="o">/</span><span class="n">jax</span><span class="o">/</span><span class="n">_src</span><span class="o">/</span><span class="n">numpy</span><span class="o">/</span><span class="n">fft</span><span class="o">.</span><span class="n">py</span><span class="p">:</span><span class="mi">19</span><span class="p">,</span> <span class="ow">in</span> <span class="o">&lt;</span><span class="n">module</span><span class="o">&gt;</span>
<span class="g g-Whitespace"> </span><span class="mi">16</span> <span class="kn">import</span> <span class="nn">operator</span>
<span class="g g-Whitespace"> </span><span class="mi">17</span> <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="ne">---&gt; </span><span class="mi">19</span> <span class="kn">from</span> <span class="nn">jax</span> <span class="kn">import</span> <span class="n">lax</span>
<span class="g g-Whitespace"> </span><span class="mi">20</span> <span class="kn">from</span> <span class="nn">jax._src.lib</span> <span class="kn">import</span> <span class="n">xla_client</span>
<span class="g g-Whitespace"> </span><span class="mi">21</span> <span class="kn">from</span> <span class="nn">jax._src.util</span> <span class="kn">import</span> <span class="n">safe_zip</span>
<span class="n">File</span> <span class="o">~/</span><span class="n">miniforge3</span><span class="o">/</span><span class="n">envs</span><span class="o">/</span><span class="n">myenv</span><span class="o">/</span><span class="n">lib</span><span class="o">/</span><span class="n">python3</span><span class="mf">.9</span><span class="o">/</span><span class="n">site</span><span class="o">-</span><span class="n">packages</span><span class="o">/</span><span class="n">jax</span><span class="o">/</span><span class="n">lax</span><span class="o">/</span><span class="fm">__init__</span><span class="o">.</span><span class="n">py</span><span class="p">:</span><span class="mi">332</span><span class="p">,</span> <span class="ow">in</span> <span class="o">&lt;</span><span class="n">module</span><span class="o">&gt;</span>
<span class="g g-Whitespace"> </span><span class="mi">299</span> <span class="kn">from</span> <span class="nn">jax._src.lax.lax</span> <span class="kn">import</span> <span class="p">(</span><span class="n">_reduce_sum</span><span class="p">,</span> <span class="n">_reduce_max</span><span class="p">,</span> <span class="n">_reduce_min</span><span class="p">,</span> <span class="n">_reduce_or</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">300</span> <span class="n">_reduce_and</span><span class="p">,</span> <span class="n">_reduce_window_sum</span><span class="p">,</span> <span class="n">_reduce_window_max</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">301</span> <span class="n">_reduce_window_min</span><span class="p">,</span> <span class="n">_reduce_window_prod</span><span class="p">,</span>
<span class="p">(</span><span class="o">...</span><span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">306</span> <span class="n">_upcast_fp16_for_computation</span><span class="p">,</span> <span class="n">_broadcasting_shape_rule</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">307</span> <span class="n">_eye</span><span class="p">,</span> <span class="n">_tri</span><span class="p">,</span> <span class="n">_delta</span><span class="p">,</span> <span class="n">_ones</span><span class="p">,</span> <span class="n">_zeros</span><span class="p">,</span> <span class="n">_dilate_shape</span><span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">308</span> <span class="kn">from</span> <span class="nn">jax._src.lax.control_flow</span> <span class="kn">import</span> <span class="p">(</span>
<span class="g g-Whitespace"> </span><span class="mi">309</span> <span class="n">associative_scan</span> <span class="k">as</span> <span class="n">associative_scan</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">310</span> <span class="n">cond</span> <span class="k">as</span> <span class="n">cond</span><span class="p">,</span>
<span class="p">(</span><span class="o">...</span><span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">330</span> <span class="n">while_p</span> <span class="k">as</span> <span class="n">while_p</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">331</span> <span class="p">)</span>
<span class="ne">--&gt; </span><span class="mi">332</span> <span class="kn">from</span> <span class="nn">jax._src.lax.fft</span> <span class="kn">import</span> <span class="p">(</span>
<span class="g g-Whitespace"> </span><span class="mi">333</span> <span class="n">fft</span> <span class="k">as</span> <span class="n">fft</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">334</span> <span class="n">fft_p</span> <span class="k">as</span> <span class="n">fft_p</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">335</span> <span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">336</span> <span class="kn">from</span> <span class="nn">jax._src.lax.parallel</span> <span class="kn">import</span> <span class="p">(</span>
<span class="g g-Whitespace"> </span><span class="mi">337</span> <span class="n">all_gather</span> <span class="k">as</span> <span class="n">all_gather</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">338</span> <span class="n">all_to_all</span> <span class="k">as</span> <span class="n">all_to_all</span><span class="p">,</span>
<span class="p">(</span><span class="o">...</span><span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">355</span> <span class="n">xeinsum</span> <span class="k">as</span> <span class="n">xeinsum</span><span class="p">,</span>
<span class="g g-Whitespace"> </span><span class="mi">356</span> <span class="p">)</span>
<span class="g g-Whitespace"> </span><span class="mi">357</span> <span class="kn">from</span> <span class="nn">jax._src.lax.other</span> <span class="kn">import</span> <span class="p">(</span>
<span class="g g-Whitespace"> </span><span class="mi">358</span> <span class="n">conv_general_dilated_patches</span> <span class="k">as</span> <span class="n">conv_general_dilated_patches</span>
<span class="g g-Whitespace"> </span><span class="mi">359</span> <span class="p">)</span>
<span class="n">File</span> <span class="o">~/</span><span class="n">miniforge3</span><span class="o">/</span><span class="n">envs</span><span class="o">/</span><span class="n">myenv</span><span class="o">/</span><span class="n">lib</span><span class="o">/</span><span class="n">python3</span><span class="mf">.9</span><span class="o">/</span><span class="n">site</span><span class="o">-</span><span class="n">packages</span><span class="o">/</span><span class="n">jax</span><span class="o">/</span><span class="n">_src</span><span class="o">/</span><span class="n">lax</span><span class="o">/</span><span class="n">fft</span><span class="o">.</span><span class="n">py</span><span class="p">:</span><span class="mi">145</span><span class="p">,</span> <span class="ow">in</span> <span class="o">&lt;</span><span class="n">module</span><span class="o">&gt;</span>
<span class="g g-Whitespace"> </span><span class="mi">143</span> <span class="n">batching</span><span class="o">.</span><span class="n">primitive_batchers</span><span class="p">[</span><span class="n">fft_p</span><span class="p">]</span> <span class="o">=</span> <span class="n">fft_batching_rule</span>
<span class="g g-Whitespace"> </span><span class="mi">144</span> <span class="k">if</span> <span class="n">pocketfft</span><span class="p">:</span>
<span class="ne">--&gt; </span><span class="mi">145</span> <span class="n">xla</span><span class="o">.</span><span class="n">backend_specific_translations</span><span class="p">[</span><span class="s1">&#39;cpu&#39;</span><span class="p">][</span><span class="n">fft_p</span><span class="p">]</span> <span class="o">=</span> <span class="n">pocketfft</span><span class="o">.</span><span class="n">pocketfft</span>
<span class="ne">AttributeError</span>: module &#39;jaxlib.pocketfft&#39; has no attribute &#39;pocketfft&#39;
</pre></div>
</div>
</div>