[CLI]: jax.Array not serializable and not working with wandb.log

[weiyaw]

Describe the bug

Jax has recently introduce a new Array type, but wandb.log no longer works with this array.

This problem will go away after setting jax.config.update('jax_array', False) but this setting is only guaranteed to exist until 15/3/2023.

Some background info:
https://jax.readthedocs.io/en/latest/jax_array_migration.html#jax-array-migration

import jax.numpy as jnp
import wandb
wandb.init()
wandb.log({'test': jnp.array([3])})

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/site-packages/wandb/sdk/wandb_run.py", line 292, in wrapper
    return func(self, *args, **kwargs)
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/site-packages/wandb/sdk/wandb_run.py", line 255, in wrapper
    return func(self, *args, **kwargs)
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/site-packages/wandb/sdk/wandb_run.py", line 1612, in log
    self._log(data=data, step=step, commit=commit)
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/site-packages/wandb/sdk/wandb_run.py", line 1394, in _log
    self._partial_history_callback(data, step, commit)
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/site-packages/wandb/sdk/wandb_run.py", line 1278, in _partial_history_callback
    self._backend.interface.publish_partial_history(
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/site-packages/wandb/sdk/interface/interface.py", line 556, in publish_partial_history
    item.value_json = json_dumps_safer_history(v)
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/site-packages/wandb/util.py", line 920, in json_dumps_safer_history
    return json.dumps(obj, cls=WandBHistoryJSONEncoder, **kwargs)
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/json/__init__.py", line 234, in dumps
    return cls(
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/site-packages/wandb/util.py", line 887, in default
    return json.JSONEncoder.default(self, obj)
  File "/Users/kenyonn/miniconda3/envs/jax4/lib/python3.9/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type Array is not JSON serializable

Additional Files

No response

Environment

WandB version: 0.13.7

OS: MacOS 12.6.1

Python version: 3.9.13

Versions of relevant libraries: jax 0.4.1

Additional Context

No response

[kevinzakka]

+1, seeing this issue too.

[ramit-wandb]

Hi @weiyaw,

Thanks for pointing this out. I'll make an internal ticket to track this request. Since Array is a new datatype within JAX, it was not supported out of the box.

One of the members from our team will reach out once this is resolved.

Thanks,
Ramit

[sangtanisakshi]

+1, Also getting the same error on more or less a similar environment.

[dmitryduev]

Hey folks, the fix for this (#4718) is in main and will be released this week, but in the meantime, would you mind checking that things work for you if installing from main?

[weiyaw]

Thanks @dmitryduev. It's working on my machine.

[kptkin]

closing since this is fixed with release 0.13.8 if someone experience another issues, please feel free to re-open.

[exalate-issue-sync]

WandB Internal User commented:
sangtanisakshi commented:
+1, Also getting the same error on more or less a similar environment.

[exalate-issue-sync]

WandB Internal User commented:
weiyaw commented:
Thanks @dmitryduev. It's working on my machine.

[exalate-issue-sync]

WandB Internal User commented:
kptkin commented:
closing since this is fixed with release 0.13.8 if someone experience another issues, please feel free to re-open.