<h3 dir="auto">Description</h3>
<ul dir="auto">
<li>
<p dir="auto">Hello! I'm getting some bad behavior when using lax.scan to compute an exponential moving average over a time axis.</p>
</li>
<li>
<p dir="auto">Specifically, the EMA function here is causing issues <a href="https://github.com/google-research/chirp/blob/main/chirp/audio_utils.py#L37">Here</a></p>
</li>
<li>
<p dir="auto">The model exports to TF SavedModel (with enable_xla=False), but then we get an exception when running the model:</p>
</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[google3/third_party/tensorflow/python/eager/execute.py](https://colab.corp.google.com/drive/1FX99EPcaX-1mAVnpnpUQwkkR0WZ_6o3h#) in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     52   try:
     53     ctx.ensure_initialized()
---&gt; 54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     55                                         inputs, attrs, num_outputs)
     56   except core._NotOkStatusException as e:

InvalidArgumentError: Must have updates.shape = indices.shape[:batch_dim] + buffer_shape[num_index_dims:], got updates.shape: [1,1,1,16], indices.shape: [1,128,1,16,1], buffer_shape: [2048], num_index_dims: 1, and batch_dim: 4

	 [[{{node jax2tf_infer_fn_/scan/while/body/ScatterNd_1}}]]

	 [[PartitionedCall/jax2tf_infer_fn_/scan/while]] [Op:__inference_restored_function_body_8540]
"><pre class="notranslate"><code class="notranslate">[google3/third_party/tensorflow/python/eager/execute.py](https://colab.corp.google.com/drive/1FX99EPcaX-1mAVnpnpUQwkkR0WZ_6o3h#) in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     52   try:
     53     ctx.ensure_initialized()
---&gt; 54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     55                                         inputs, attrs, num_outputs)
     56   except core._NotOkStatusException as e:

InvalidArgumentError: Must have updates.shape = indices.shape[:batch_dim] + buffer_shape[num_index_dims:], got updates.shape: [1,1,1,16], indices.shape: [1,128,1,16,1], buffer_shape: [2048], num_index_dims: 1, and batch_dim: 4

	 [[{{node jax2tf_infer_fn_/scan/while/body/ScatterNd_1}}]]

	 [[PartitionedCall/jax2tf_infer_fn_/scan/while]] [Op:__inference_restored_function_body_8540]

</code></pre></div>
<p dir="auto">which suggests some indices are being misplaced.</p>
<ul dir="auto">
<li>
<p dir="auto">A couple observations:</p>
</li>
<li>
<p dir="auto">The error occurs even when polymorphic shape is disabled.</p>
</li>
<li>
<p dir="auto">When xla is enabled, the savedmodel works correctly.</p>
</li>
<li>
<p dir="auto">Strangely, converting the SavedModel to TFLite seems to work out; the model runs and the outptus match the jax model numerically.</p>
</li>
<li>
<p dir="auto">Minimal reproducing code is in my usual colab notebook of minimal reproductions <a href="https://colab.corp.google.com/drive/1FX99EPcaX-1mAVnpnpUQwkkR0WZ_6o3h#scrollTo=280PugPMoiQH" rel="nofollow">Here</a></p>
</li>
</ul>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">v0.3.17</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">CPU, GPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Linux</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<p dir="auto"><em>No response</em></p>
<h3 dir="auto">Description</h3>
<p dir="auto">Hello! I'm getting some bad behavior when using lax.scan to compute an exponential moving average over a time axis.</p>
<p dir="auto">Specifically, the EMA function here is causing issues:<br>
<a href="https://github.com/google-research/chirp/blob/main/chirp/audio_utils.py#L37">https://github.com/google-research/chirp/blob/main/chirp/audio_utils.py#L37</a></p>
<p dir="auto">The model exports to TF SavedModel (with enable_xla=False), but then we get an exception when running the model:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[google3/third_party/tensorflow/python/eager/execute.py](https://colab.corp.google.com/drive/1FX99EPcaX-1mAVnpnpUQwkkR0WZ_6o3h#) in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     52   try:
     53     ctx.ensure_initialized()
---&gt; 54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     55                                         inputs, attrs, num_outputs)
     56   except core._NotOkStatusException as e:

InvalidArgumentError: Must have updates.shape = indices.shape[:batch_dim] + buffer_shape[num_index_dims:], got updates.shape: [1,1,1,16], indices.shape: [1,128,1,16,1], buffer_shape: [2048], num_index_dims: 1, and batch_dim: 4

	 [[{{node jax2tf_infer_fn_/scan/while/body/ScatterNd_1}}]]

	 [[PartitionedCall/jax2tf_infer_fn_/scan/while]] [Op:__inference_restored_function_body_8540]"><pre class="notranslate"><code class="notranslate">[google3/third_party/tensorflow/python/eager/execute.py](https://colab.corp.google.com/drive/1FX99EPcaX-1mAVnpnpUQwkkR0WZ_6o3h#) in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     52   try:
     53     ctx.ensure_initialized()
---&gt; 54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     55                                         inputs, attrs, num_outputs)
     56   except core._NotOkStatusException as e:

InvalidArgumentError: Must have updates.shape = indices.shape[:batch_dim] + buffer_shape[num_index_dims:], got updates.shape: [1,1,1,16], indices.shape: [1,128,1,16,1], buffer_shape: [2048], num_index_dims: 1, and batch_dim: 4

	 [[{{node jax2tf_infer_fn_/scan/while/body/ScatterNd_1}}]]

	 [[PartitionedCall/jax2tf_infer_fn_/scan/while]] [Op:__inference_restored_function_body_8540]
</code></pre></div>
<p dir="auto">which suggests some indices are being misplaced.</p>
<p dir="auto">A couple observations:</p>
<ul dir="auto">
<li>The error occurs even when polymorphic shape is disabled.</li>
<li>When xla is enabled, the savedmodel works correctly.</li>
<li>Strangely, converting the SavedModel to TFLite seems to work out; the model runs and the outptus match the jax model numerically.</li>
</ul>
<p dir="auto">Minimal reproducing code is in my usual colab notebook of minimal reproductions:<br>
<a href="https://colab.corp.google.com/drive/1FX99EPcaX-1mAVnpnpUQwkkR0WZ_6o3h#scrollTo=280PugPMoiQH" rel="nofollow">https://colab.corp.google.com/drive/1FX99EPcaX-1mAVnpnpUQwkkR0WZ_6o3h#scrollTo=280PugPMoiQH</a></p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">v0.3.17</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">CPU, GPU</p>
<h3 dir="auto">Additional System Info</h3>
<p dir="auto">Linux</p>