<p dir="auto"><code class="notranslate">~/jax$ python3 tests/host_callback_test.py</code> fails with:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[ RUN      ] CallJaxTest.test_jax_grad
2021-04-15 00:16:53.333864: E external/org_tensorflow/tensorflow/compiler/xla/status_macros.cc:56] Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape u32[2]{0} and outfeed source buffer shape f32[]"><pre class="notranslate"><code class="notranslate">[ RUN      ] CallJaxTest.test_jax_grad
2021-04-15 00:16:53.333864: E external/org_tensorflow/tensorflow/compiler/xla/status_macros.cc:56] Internal: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/service/gpu/outfeed_thunk.cc:80) ShapeUtil::Equal(source_slices_[index].shape, output_shape) Mismatch between outfeed output buffer shape u32[2]{0} and outfeed source buffer shape f32[]
</code></pre></div>
<p dir="auto"><code class="notranslate">Mismatch between outfeed output buffer shape u32[2]{0} and outfeed source buffer shape f32[]</code></p>
<p dir="auto">All test cases seem to fail this way. It's a little hard to tell because after a few test cases run and fail with this, the test will just hang until you kill it (probably because computations are blocked waiting for infeed).</p>
<p dir="auto">I haven't dug into this at all, but for now I'm planning to disable all of host_callback_test on multi-GPU platforms in order to make our test suite runnable on multi-GPU (since right now it hangs).</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 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">Hi all,</p>
<p dir="auto">Whilst working on a recent project I realised that scipy.stats.multinomial doesn't have a jax.scipy implementation.</p>
<p dir="auto">I thought it'd be worthwhile to implement as well as a decent first contribution.</p>
<p dir="auto">I think I'm largely done on the work <a href="https://github.com/harryjulian/jax/tree/scipy.stats.multinomial">here</a> but thought to make an issue before making a PR.</p>
<p dir="auto">Is there any reason this hasn't been implemented thus far?</p>