Improve how Jax array is handled

[sunnyayyl]

import jax.numpy as jnp
import lets_plot as lp
import numpy

(
    lp.ggplot(
        {"x": list(jnp.arange(0, 10)), "y": numpy.arange(0, 10)},
        lp.aes(x="x", y="y"),
    )

    + lp.geom_line()
).show()

y axis is the expected behavior with numpy x axis is the behavior when used with jax (jax: https://jax.readthedocs.io/en/latest/index.html)

[alshan]

Hi, in the example above, you are passing a list of arrays. When the elements in the list are anything but numeric values, Lets-Plot treats them as discrete (a.k.a. categorical) values and creates a discrete axis where the labels are simple string representations of the elements in the list.

A bare jax array, however, will most likely not be recognized as a collection of numerics either. Did you try this?

[sunnyayyl]

I know I can use numpy.asarray, but it would be nice if the conversion can be done implicitly, so I guess this is more of a feature request.

[alshan]

Sure, thanks for the feature request.

[sunnyayyl]

I'll edit the title to better reflect that this issue is a feature request. Thanks for your help!

[sunnyayyl]

I had a look at the source code, would this issue be solved by adding to this part? https://github.com/JetBrains/lets-plot/blob/master/python-package%2Flets_plot%2F_type_utils.py

[sunnyayyl]

Something on the line of

try:
   import jax.numpy as jnp
except:
   jnp=None

if (jnp and isinstance(v, jnp.ndarray))
   return _standardize_value(v.tolist())

(Edit: annoyingly, I will not be next to a computer for the next few days, so I can't test this)

[alshan]

This looks promising, thanx!

[sunnyayyl]

I'm currently working on this, I will try and open a pull request when I am done

[sunnyayyl]