[JAX] The use of jax.Array.at fails experimental_from_jax
To fix this, the expected output of the **_update** function should be printed before and after the interpreter is invoked, and the output of the interpreter should be captured and printed. This can be done using the following code:

```
args = jnp.ones(2)

expected = _update(args)
print("Expected output:", expected)

interpreter.set_tensor(input_details[0]['index'], args)
interpreter.invoke()

output = interpreter.get_tensor(output_details[0]['index'])
print("Output:", output)
``` > To fix this, the expected output of the **_update** function should be printed before and after the interpreter is invoked, and the output of the interpreter should be captured and printed. This can be done using the following code:
> 
> ```
> args = jnp.ones(2)
> 
> expected = _update(args)
> print("Expected output:", expected)
> 
> interpreter.set_tensor(input_details[0]['index'], args)
> interpreter.invoke()
> 
> output = interpreter.get_tensor(output_details[0]['index'])
> print("Output:", output)
> ```

Nope, still got the same  error  (the same happens for tf 2.12)

https://colab.research.google.com/drive/1VGVrdBBqN_0qrvZC_w99ZxVTJqp6_9Nt?usp=sharing
 PS

The same functionality can one implemented with `jax.nn.one_hot` and this works. So clearly the issue is with at.

@tiruk007 does it make sense to try the path Jax->jax2tf->tf.lite or is this under the hood of `tf.lite.TFLiteConverter.experimental_from_jax` ?
 @krzysztofrusek 
Sorry for the late reply.
@pjpratik 
Could you please look into this? I was able to reproduce the issue on Colab using TFv2.12. Please find the gist [here](https://colab.research.google.com/gist/tiruk007/09d4f2106e8658d7ad5cc23314706a7c/untitled169.ipynb) for reference.

Thank you ! Hi @krzysztofrusek

The error occurs while we are updating a scalar value and as the error suggests when we use `jp.ones((1,2))` and try to update `x.at[0].set[4]`, it runs without any error.  Please find the [gist](https://colab.research.google.com/gist/pjpratik/0ef750b5a7b394c700ad42770ebac4a3/60090.ipynb) here.  

@sachinprasadhs Could you please check if this is intended behaviour? 

Thanks. Hi @pjpratik thanks for the hint. This is some workaround. Hi @krzysztofrusek 

The latest nightly TF 2.14.0-dev20230716 says
```
TFLiteConverterV2.experimental_from_jax (from tensorflow.lite.python.lite) is deprecated and will be removed in a future version.
Instructions for updating:
Use `jax2tf.convert` and (`lite.TFLiteConverter.from_saved_model` or `lite.TFLiteConverter.from_concrete_functions`) instead.
```
So the right way is to jax->jax2tf->lite moving forward.

Please find the resolved [gist](https://colab.research.google.com/gist/pjpratik/b9c36fbf4630d220768454e6262cc249/60090.ipynb) for the same.

Thanks.  This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60090">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60090">No</a>

### 1. System information

- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Darwin Kernel Version 22.3.0
- TensorFlow installation (pip package or built from source): pip
- TensorFlow library (version, if pip package or github SHA, if built from source): 2.11.0 

### 2. Code


#### Option A: Reference colab notebooks

1)  Reference [TensorFlow Model Colab](https://colab.research.google.com/gist/ymodak/e96a4270b953201d5362c61c1e8b78aa/tensorflow-datasets.ipynb?authuser=1): Demonstrate how to build your TF model.
2)  Reference [TensorFlow Lite Model Colab](https://colab.research.google.com/gist/ymodak/0dfeb28255e189c5c48d9093f296e9a8/tensorflow-lite-debugger-colab.ipynb): Demonstrate how to convert your TF model to a TF Lite model (with quantization, if used) and run TFLite Inference (if possible).

```
(You can paste links or attach files by dragging & dropping them below)
- Provide links to your updated versions of the above two colab notebooks.
- Provide links to your TensorFlow model and (optionally) TensorFlow Lite Model.
```

#### Option B: Paste your code here or provide a link to a custom end-to-end colab

```

import tensorflow as tf
import jax
import jax.numpy as jnp
import jax.tree_util as tree

def main3():
    print(f'JAX {jax.__version__}')
    print(f'tf {tf.__version__}')

    @jax.jit
    def _update(x):
        return x.at[0].set(4)


    converter = tf.lite.TFLiteConverter.experimental_from_jax([_update],
                                                              [[('x', jnp.ones(2))]])
    converter.target_spec.supported_ops = [
        tf.lite.OpsSet.TFLITE_BUILTINS,  # enable TensorFlow Lite ops.
        tf.lite.OpsSet.SELECT_TF_OPS  # enable TensorFlow ops.

    ]

    tflite_update = converter.convert()
    with open('update.tflite', 'wb') as f:
        f.write(tflite_update)

    interpreter = tf.lite.Interpreter(model_content=tflite_update)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()



    args = jnp.ones(2)

    expected = _update(args)

    # for a, d in zip(args, input_details):
    interpreter.set_tensor(input_details[0]['index'], args)

    interpreter.invoke()
    return

```

full [example](https://github.com/krzysztofrusek/reinforced-lib/blob/lite/examples/lite/lite.py)

### 3. Failure after conversion
If the conversion is successful, but the generated model is wrong, then state what is wrong:

- Conversion works
- Interpreter fails at invoke

### 4. (optional) RNN conversion support
If converting TF RNN to TFLite fused RNN ops, please prefix [RNN] in the title.

### 5. (optional) Any other info / logs
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.

```Traceback (most recent call last):
  File "/Users/krzysiek/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/223.8836.34/PyCharm.app/Contents/plugins/python/helpers/pydev/pydevd.py", line 1496, in _exec
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "/Users/krzysiek/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/223.8836.34/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/Users/krzysiek/Documents/ml4wifi/reinforced-lib/examples/lite/lite.py", line 168, in <module>
    main3()
  File "/Users/krzysiek/Documents/ml4wifi/reinforced-lib/examples/lite/lite.py", line 163, in main3
    interpreter.invoke()
  File "/Users/krzysiek/Documents/ml4wifi/ftmrate_internal/venv/lib/python3.10/site-packages/tensorflow/lite/python/interpreter.py", line 917, in invoke
    self._interpreter.Invoke()
RuntimeError: Updates shape must have rank at least one. Found:[]Node number 1 (TfLiteFlexDelegate) failed to invoke.
```
