Merge pull request #29 from blab/switching-to-jax.Array

Updating Array types

evofr/models/piantham_model.py

@@ -28,7 +28,7 @@ def compute_frequency_piantham(ga, q0, gen_rev, T):
 
     Returns
     -------
-    Simulated frequencies as DeviceArray.
+    Simulated frequencies as Array.
     """
     _ga = jnp.append(ga, 1.0)
     max_age = gen_rev.shape[-1]

evofr/plotting/plot_functions.py

@@ -17,7 +17,7 @@ def prep_posterior_for_plot(
 
     samples:
         Dictionary with keys being site or variable names.
-        Values are DeviceArrays containing posterior samples
+        Values are Arrays containing posterior samples
         with shape (sample_number, site_shape).
 
     ps:
@@ -65,7 +65,7 @@ def plot_posterior_time(
         Median values.
 
     quants:
-        Quantiles to be plotted. Organized as a list of CIs as DeviceArrays.
+        Quantiles to be plotted. Organized as a list of CIs as Arrays.
 
     alphas:
         Transparency for each quantile.

evofr/posterior/posterior_handler.py

@@ -18,7 +18,7 @@ def __init__(
         ----------
         samples:
             optional dictionary with keys corresponding to variable names.
-            Values will be DeviceArrays containing posterior samples.
+            Values will be Arrays containing posterior samples.
 
         data:
             optional DataSpec instance containing underlying data from analysis

evofr/posterior/posterior_helpers.py

@@ -1,5 +1,6 @@
 from functools import partial
 from typing import Dict, List, Optional
+import jax
 import jax.numpy as jnp
 import json
 import numpy as np
@@ -16,7 +17,7 @@ def get_quantile(samples: Dict, p, site):
     ----------
     samples:
         Dictionary with keys being site or variable names.
-        Values are DeviceArrays with shape (sample_number, site_shape).
+        Values are Arrays with shape (sample_number, site_shape).
 
     p:
         Percent credible interval to return.
@@ -26,7 +27,7 @@ def get_quantile(samples: Dict, p, site):
 
     Returns
     -------
-    DeviceArray of shape (site_shape).
+    Array of shape (site_shape).
     """
     q = jnp.array([0.5 * (1 - p), 0.5 * (1 + p)])
     return jnp.quantile(samples[site], q=q, axis=0)
@@ -156,7 +157,7 @@ def default(self, obj):
             return round(float(obj), 3)
         if isinstance(obj, np.ndarray):
             return obj.tolist()
-        if isinstance(obj, jnp.DeviceArray):
+        if isinstance(obj, jax.Array):
             return self.default(np.array(obj))
         if isinstance(obj, pd.Timestamp):
             return obj.strftime("%Y-%m-%d")