<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="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; import jax.numpy as jnp
&gt;&gt;&gt; a = [1, 3, 3, 5]
&gt;&gt;&gt; t = jnp.asarray(a, dtype=jnp.bfloat16)
&gt;&gt;&gt; import pickle
&gt;&gt;&gt; pp = pickle.dump(t, file=open('test.pkl','wb'))
Traceback (most recent call last):
  File &quot;&lt;stdin&gt;&quot;, line 1, in &lt;module&gt;
_pickle.PicklingError: Can't pickle &lt;class 'bfloat16'&gt;: attribute lookup bfloat16 on builtins failed

&gt;&gt;&gt; import jax
&gt;&gt;&gt; jax.__version__
'0.2.25'
&gt;&gt;&gt; import jaxlib
&gt;&gt;&gt; jaxlib.__version__
'0.1.73'"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</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-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">a</span> <span class="pl-c1">=</span> [<span class="pl-c1">1</span>, <span class="pl-c1">3</span>, <span class="pl-c1">3</span>, <span class="pl-c1">5</span>]
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">t</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">asarray</span>(<span class="pl-s1">a</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-s1">bfloat16</span>)
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">import</span> <span class="pl-s1">pickle</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">pp</span> <span class="pl-c1">=</span> <span class="pl-s1">pickle</span>.<span class="pl-en">dump</span>(<span class="pl-s1">t</span>, <span class="pl-s1">file</span><span class="pl-c1">=</span><span class="pl-en">open</span>(<span class="pl-s">'test.pkl'</span>,<span class="pl-s">'wb'</span>))
<span class="pl-v">Traceback</span> (<span class="pl-s1">most</span> <span class="pl-s1">recent</span> <span class="pl-s1">call</span> <span class="pl-s1">last</span>):
  <span class="pl-v">File</span> <span class="pl-s">"&lt;stdin&gt;"</span>, <span class="pl-s1">line</span> <span class="pl-c1">1</span>, <span class="pl-s1">in</span> <span class="pl-c1">&lt;</span><span class="pl-s1">module</span><span class="pl-c1">&gt;</span>
<span class="pl-s1">_pickle</span>.<span class="pl-v">PicklingError</span>: <span class="pl-v">Can</span><span class="pl-s">'t pickle &lt;class '</span><span class="pl-s1">bfloat16</span>'<span class="pl-c1">&gt;</span>: <span class="pl-s1">attribute</span> <span class="pl-s1">lookup</span> <span class="pl-s1">bfloat16</span> <span class="pl-s1">on</span> <span class="pl-s1">builtins</span> <span class="pl-s1">failed</span>

<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">jax</span>.<span class="pl-s1">__version__</span>
<span class="pl-s">'0.2.25'</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">import</span> <span class="pl-s1">jaxlib</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">jaxlib</span>.<span class="pl-s1">__version__</span>
<span class="pl-s">'0.1.73'</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">Jax2tf currently does not support scatter_* in non XLA terms (see <a href="https://github.com/google/jax/blob/main/jax/experimental/jax2tf/converters_eval/no_xla_limitations.md">docs</a>)</p>
<p dir="auto">I am planning on implementing this as I need it, but first I wanted to check - is it possible, or is there a very solid reason for why it simply isn't expressible in the tensorflow ops?</p>
<p dir="auto">Is this a foolish plan?</p>
<p dir="auto">I want to do it in non-XLA terms so I can use tf2onnx so it can be used as an ONNX model</p>
<p dir="auto">I was thinking about using a combination of:</p>
<ul dir="auto">
<li><code class="notranslate">tf.math.unsorted_segment_{prod,sum,...}</code> for getting the updates to the array</li>
<li><code class="notranslate">tf.scatter_nd</code> for applying those updates</li>
</ul>
<p dir="auto">Note (to self): <code class="notranslate">x.at[indexs].get()</code> does has a non-XLA implementation.</p>