Issue 1:
Description

Hello! I'm getting some bad behavior when using lax.scan to compute an exponential moving average over a time axis.

Specifically, the EMA function here is causing issues Here

The model exports to TF SavedModel (with enable_xla=False), but then we get an exception when running the model:

[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()
---> 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]

which suggests some indices are being misplaced.

A couple observations:

The error occurs even when polymorphic shape is disabled.

When xla is enabled, the savedmodel works correctly.

Strangely, converting the SavedModel to TFLite seems to work out; the model runs and the outptus match the jax model numerically.

Minimal reproducing code is in my usual colab notebook of minimal reproductions Here

What jax/jaxlib version are you using?

v0.3.17

Which accelerator(s) are you using?

CPU, GPU

Additional system info

Linux

NVIDIA GPU info

No response

Issue 2:
Description

Hello! I'm getting some bad behavior when using lax.scan to compute an exponential moving average over a time axis.

Specifically, the EMA function here is causing issues:

https://github.com/google-research/chirp/blob/main/chirp/audio_utils.py#L37

The model exports to TF SavedModel (with enable_xla=False), but then we get an exception when running the model:

[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()
---> 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]

which suggests some indices are being misplaced.

A couple observations:

The error occurs even when polymorphic shape is disabled.

When xla is enabled, the savedmodel works correctly.

Strangely, converting the SavedModel to TFLite seems to work out; the model runs and the outptus match the jax model numerically.

Minimal reproducing code is in my usual colab notebook of minimal reproductions:

https://colab.corp.google.com/drive/1FX99EPcaX-1mAVnpnpUQwkkR0WZ_6o3h#scrollTo=280PugPMoiQH

What jax/jaxlib version are you using?

v0.3.17

Which accelerator(s) are you using?

CPU, GPU

Additional System Info

Linux