Deploying to gh-pages from @ a5cb116 🚀

Author: lucascolley

generated/array_api_extra.at.html

@@ -367,6 +367,24 @@ <h1>array_api_extra.at<a class="headerlink" href="#array-api-extra-at" title="Li
 <p><a class="reference external" href="https://sparse.pydata.org/">sparse</a>, as well as read-only arrays from libraries
 not explicitly covered by <code class="docutils literal notranslate"><span class="pre">array-api-compat</span></code>, are not supported by update
 methods.</p>
+<p>Boolean masks are supported on Dask and jitted JAX arrays exclusively
+when <cite>idx</cite> has the same shape as <cite>x</cite> and <cite>y</cite> is 0-dimensional.
+Note that this support is not available in JAX’s native
+<code class="docutils literal notranslate"><span class="pre">x.at[mask].set(y)</span></code>.</p>
+<p>This pattern:</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">mask</span> <span class="o">=</span> <span class="n">m</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">x</span><span class="p">[</span><span class="n">mask</span><span class="p">]</span> <span class="o">=</span> <span class="n">f</span><span class="p">(</span><span class="n">x</span><span class="p">[</span><span class="n">mask</span><span class="p">])</span>
+</pre></div>
+</div>
+<p>Can’t be replaced by <cite>at</cite>, as it won’t work on Dask and JAX inside jax.jit:</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">mask</span> <span class="o">=</span> <span class="n">m</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">x</span> <span class="o">=</span> <span class="n">xpx</span><span class="o">.</span><span class="n">at</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">mask</span><span class="p">)</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="n">f</span><span class="p">(</span><span class="n">x</span><span class="p">[</span><span class="n">mask</span><span class="p">])</span>  <span class="c1"># Crash on Dask and jax.jit</span>
+</pre></div>
+</div>
+<p>You should instead use:</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="n">x</span> <span class="o">=</span> <span class="n">xp</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">m</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">f</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">x</span><span class="p">)</span>
+</pre></div>
+</div>
 <p class="rubric">Examples</p>
 <p>Given either of these equivalent expressions:</p>
 <div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span><span class="w"> </span><span class="nn">array_api_extra</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">xpx</span>