Update support to jax (#585)

The change is minimal since jax tries to mimic numpy already.
lanpa · Jun 28, 2020 · 15f4537 · 15f4537
1 parent 8f50866
commit 15f4537
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensorboardX/x2num.py b/tensorboardX/x2num.py
@@ -30,6 +30,8 @@ def make_np(x):
         return check_nan(prepare_chainer(x))
     if 'mxnet' in str(type(x)):
         return check_nan(prepare_mxnet(x))
+    if 'jax' in str(type(x)):
+        return check_nan(np.array(x))
     raise NotImplementedError(
         'Got {}, but expected numpy array or torch tensor.'.format(type(x)))