<p dir="auto">The code blow returns <code class="notranslate">NotImplementedError: sparse rule for xor</code>.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
from jax.experimental import sparse
indices = [[0, 1], [1, 2]]
X = sparse.BCOO(([1 for _ in range(len(indices))], indices), shape=(3, 4))
y = jnp.ones(3, dtype=bool)

@sparse.sparsify
def xor(X, y):
  return y[:, None] ^ X

xor(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">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">experimental</span> <span class="pl-k">import</span> <span class="pl-s1">sparse</span>
<span class="pl-s1">indices</span> <span class="pl-c1">=</span> [[<span class="pl-c1">0</span>, <span class="pl-c1">1</span>], [<span class="pl-c1">1</span>, <span class="pl-c1">2</span>]]
<span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-s1">sparse</span>.<span class="pl-v">BCOO</span>(([<span class="pl-c1">1</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-en">len</span>(<span class="pl-s1">indices</span>))], <span class="pl-s1">indices</span>), <span class="pl-s1">shape</span><span class="pl-c1">=</span>(<span class="pl-c1">3</span>, <span class="pl-c1">4</span>))
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-c1">3</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">bool</span>)

<span class="pl-en">@<span class="pl-s1">sparse</span>.<span class="pl-s1">sparsify</span></span>
<span class="pl-k">def</span> <span class="pl-en">xor</span>(<span class="pl-v">X</span>, <span class="pl-s1">y</span>):
  <span class="pl-k">return</span> <span class="pl-s1">y</span>[:, <span class="pl-c1">None</span>] <span class="pl-c1">^</span> <span class="pl-v">X</span>

<span class="pl-en">xor</span>(<span class="pl-v">X</span>, <span class="pl-s1">y</span>)</pre></div>
<p dir="auto">Please: add support for XOR operation over sparse matrices.</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 requests.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Describe your goal, and if possible provide a code snippet with a motivating example.</li>
</ul>
<p dir="auto">tl;dr: Please implement np.eig in JAX, in particular its forward-mode differentiation</p>
<p dir="auto">I am working on a quantum project that is currently written in TensorFlow but would greatly benefit from JAX. The main roadblock is that the np.eig function does not have its forward-mode differentiation rule implemented in JAX.</p>
<p dir="auto">Thanks a lot!</p>