<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>
</ul>
<p dir="auto">I wanted to define a <code class="notranslate">custom_vjp</code> for a function that takes a string to indicate what operations are to be conducted. Since <code class="notranslate">str</code> is not a validate type for JAX whereas any objects that can be flatten and unflatten as a <code class="notranslate">pytree</code> is, I thought of this hacky way to define my own string object:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp
from functools import partial

OP_TO_IDX, IDX_TO_OP = {&quot;add&quot;: 0, &quot;mul&quot;: 1}, {0: &quot;add&quot;, 1: &quot;mul&quot;}

@jax.tree_util.register_pytree_node_class
class OpStr(object):
    def __init__(self, op):
        assert op in OP_TO_IDX, &quot;can only be one of the ops&quot;
        self.op = op

    def __eq__(self, other):
        return self.op == other

    def __hash__(self):
        return hash(self.op)

    def tree_flatten(self):
        return ([OP_TO_IDX[self.op]]), None

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        idx = int(children[0])
        return cls(op=IDX_TO_OP[idx])

@partial(jax.custom_vjp, nondiff_argnums=(2,))
def op(x, y, op):
    if op == &quot;add&quot;: z = x + y
    if op == &quot;mul&quot;: z = x * y
    return z

def op_fwd(x, y, op):
    cache = (x, y, op)
    if op == &quot;add&quot;: z = x + y
    if op == &quot;mul&quot;: z = x * y
    return z, cache

def op_bwd(cache, dz):
    x, y, op = cache
    if op == &quot;add&quot;: dz_dx = dz_dy = 1.0
    if op == &quot;mul&quot;: dz_dx, dz_dy = y, x
    return dz_dx, dz_dy

op.defvjp(op_fwd, op_bwd)
grad_op = jax.grad(op)
grad_op(1.0, 1.0, OpStr(&quot;add&quot;))"><pre class="notranslate"><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">jnp</span>
<span class="pl-k">from</span> <span class="pl-s1">functools</span> <span class="pl-k">import</span> <span class="pl-s1">partial</span>

<span class="pl-v">OP_TO_IDX</span>, <span class="pl-v">IDX_TO_OP</span> <span class="pl-c1">=</span> {<span class="pl-s">"add"</span>: <span class="pl-c1">0</span>, <span class="pl-s">"mul"</span>: <span class="pl-c1">1</span>}, {<span class="pl-c1">0</span>: <span class="pl-s">"add"</span>, <span class="pl-c1">1</span>: <span class="pl-s">"mul"</span>}

<span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">tree_util</span>.<span class="pl-s1">register_pytree_node_class</span></span>
<span class="pl-k">class</span> <span class="pl-v">OpStr</span>(<span class="pl-s1">object</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">op</span>):
        <span class="pl-k">assert</span> <span class="pl-s1">op</span> <span class="pl-c1">in</span> <span class="pl-v">OP_TO_IDX</span>, <span class="pl-s">"can only be one of the ops"</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">op</span> <span class="pl-c1">=</span> <span class="pl-s1">op</span>

    <span class="pl-k">def</span> <span class="pl-en">__eq__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">other</span>):
        <span class="pl-k">return</span> <span class="pl-s1">self</span>.<span class="pl-s1">op</span> <span class="pl-c1">==</span> <span class="pl-s1">other</span>

    <span class="pl-k">def</span> <span class="pl-en">__hash__</span>(<span class="pl-s1">self</span>):
        <span class="pl-k">return</span> <span class="pl-en">hash</span>(<span class="pl-s1">self</span>.<span class="pl-s1">op</span>)

    <span class="pl-k">def</span> <span class="pl-en">tree_flatten</span>(<span class="pl-s1">self</span>):
        <span class="pl-k">return</span> ([<span class="pl-v">OP_TO_IDX</span>[<span class="pl-s1">self</span>.<span class="pl-s1">op</span>]]), <span class="pl-c1">None</span>

    <span class="pl-en">@<span class="pl-s1">classmethod</span></span>
    <span class="pl-k">def</span> <span class="pl-en">tree_unflatten</span>(<span class="pl-s1">cls</span>, <span class="pl-s1">aux_data</span>, <span class="pl-s1">children</span>):
        <span class="pl-s1">idx</span> <span class="pl-c1">=</span> <span class="pl-en">int</span>(<span class="pl-s1">children</span>[<span class="pl-c1">0</span>])
        <span class="pl-k">return</span> <span class="pl-en">cls</span>(<span class="pl-s1">op</span><span class="pl-c1">=</span><span class="pl-v">IDX_TO_OP</span>[<span class="pl-s1">idx</span>])

<span class="pl-en">@<span class="pl-en">partial</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">custom_vjp</span>, <span class="pl-s1">nondiff_argnums</span><span class="pl-c1">=</span>(<span class="pl-c1">2</span>,))</span>
<span class="pl-k">def</span> <span class="pl-en">op</span>(<span class="pl-s1">x</span>, <span class="pl-s1">y</span>, <span class="pl-s1">op</span>):
    <span class="pl-k">if</span> <span class="pl-s1">op</span> <span class="pl-c1">==</span> <span class="pl-s">"add"</span>: <span class="pl-s1">z</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span> <span class="pl-c1">+</span> <span class="pl-s1">y</span>
    <span class="pl-k">if</span> <span class="pl-s1">op</span> <span class="pl-c1">==</span> <span class="pl-s">"mul"</span>: <span class="pl-s1">z</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span> <span class="pl-c1">*</span> <span class="pl-s1">y</span>
    <span class="pl-k">return</span> <span class="pl-s1">z</span>

<span class="pl-k">def</span> <span class="pl-en">op_fwd</span>(<span class="pl-s1">x</span>, <span class="pl-s1">y</span>, <span class="pl-s1">op</span>):
    <span class="pl-s1">cache</span> <span class="pl-c1">=</span> (<span class="pl-s1">x</span>, <span class="pl-s1">y</span>, <span class="pl-s1">op</span>)
    <span class="pl-k">if</span> <span class="pl-s1">op</span> <span class="pl-c1">==</span> <span class="pl-s">"add"</span>: <span class="pl-s1">z</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span> <span class="pl-c1">+</span> <span class="pl-s1">y</span>
    <span class="pl-k">if</span> <span class="pl-s1">op</span> <span class="pl-c1">==</span> <span class="pl-s">"mul"</span>: <span class="pl-s1">z</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span> <span class="pl-c1">*</span> <span class="pl-s1">y</span>
    <span class="pl-k">return</span> <span class="pl-s1">z</span>, <span class="pl-s1">cache</span>

<span class="pl-k">def</span> <span class="pl-en">op_bwd</span>(<span class="pl-s1">cache</span>, <span class="pl-s1">dz</span>):
    <span class="pl-s1">x</span>, <span class="pl-s1">y</span>, <span class="pl-s1">op</span> <span class="pl-c1">=</span> <span class="pl-s1">cache</span>
    <span class="pl-k">if</span> <span class="pl-s1">op</span> <span class="pl-c1">==</span> <span class="pl-s">"add"</span>: <span class="pl-s1">dz_dx</span> <span class="pl-c1">=</span> <span class="pl-s1">dz_dy</span> <span class="pl-c1">=</span> <span class="pl-c1">1.0</span>
    <span class="pl-k">if</span> <span class="pl-s1">op</span> <span class="pl-c1">==</span> <span class="pl-s">"mul"</span>: <span class="pl-s1">dz_dx</span>, <span class="pl-s1">dz_dy</span> <span class="pl-c1">=</span> <span class="pl-s1">y</span>, <span class="pl-s1">x</span>
    <span class="pl-k">return</span> <span class="pl-s1">dz_dx</span>, <span class="pl-s1">dz_dy</span>

<span class="pl-s1">op</span>.<span class="pl-en">defvjp</span>(<span class="pl-s1">op_fwd</span>, <span class="pl-s1">op_bwd</span>)
<span class="pl-s1">grad_op</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-s1">op</span>)
<span class="pl-en">grad_op</span>(<span class="pl-c1">1.0</span>, <span class="pl-c1">1.0</span>, <span class="pl-v">OpStr</span>(<span class="pl-s">"add"</span>))</pre></div>
<ul class="contains-task-list">
<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">But this gives me</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;/Users/wangy1/Documents/GitHub/dgl/tests/jax/test_op_str.py&quot;, line 47, in &lt;module&gt;
    grad_op(1.0, 1.0, OpStr(&quot;add&quot;))
  File &quot;/Users/wangy1/Documents/GitHub/dgl/tests/jax/test_op_str.py&quot;, line 24, in tree_unflatten
    idx = int(children[0])
TypeError: int() argument must be a string, a bytes-like object or a real number, not 'NoneType'"><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "/Users/wangy1/Documents/GitHub/dgl/tests/jax/test_op_str.py", line 47, in &lt;module&gt;
    grad_op(1.0, 1.0, OpStr("add"))
  File "/Users/wangy1/Documents/GitHub/dgl/tests/jax/test_op_str.py", line 24, in tree_unflatten
    idx = int(children[0])
TypeError: int() argument must be a string, a bytes-like object or a real number, not 'NoneType'
</code></pre></div>
<p dir="auto">Which is not expected since the flattened tree shouldn't be <code class="notranslate">None</code>.</p>
<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>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="jnp.take(
    jnp.arange(5),
    jnp.arange(5 * 3),
    axis=0
)
# DeviceArray([0, 1, 2, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4], dtype=int32)

jnp.take(
    jnp.arange(5),
    jnp.arange(5 * 3),
    axis=0,
    mode=&quot;raise&quot;,
)
# NotImplementedError"><pre class="notranslate"><code class="notranslate">jnp.take(
    jnp.arange(5),
    jnp.arange(5 * 3),
    axis=0
)
# DeviceArray([0, 1, 2, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4], dtype=int32)

jnp.take(
    jnp.arange(5),
    jnp.arange(5 * 3),
    axis=0,
    mode="raise",
)
# NotImplementedError
</code></pre></div>
<p dir="auto">According to the <a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.take.html" rel="nofollow">documentation</a>, the default out-of-bound behavior is "raise". However, this is not the case in practice</p>