parameter server strategy

distml/operator/jax_operator.py

@@ -1,5 +1,6 @@
 import numpy as np
 import cupy as cp
+
 import jax
 from jax import value_and_grad
 import jax.numpy as jnp
@@ -349,6 +350,36 @@ def reset_optimizer_for_params(self, params):
         self.tree = tree_structure(params)
         self.opt_state = self.opt_init(params)
 
+    def ones(self, shape, cpu=True):
+        if cpu:
+            return np.ones(shape)
+        else:
+            return jnp.ones(shape)
+
+    def zeros(self, shape, cpu=True):
+        if cpu:
+            return np.zeros(shape)
+        else:
+            return jnp.zeros(shape)
+
+    def ones_like(self, x, cpu=True):
+        if cpu:
+            return np.ones_like(x)
+        else:
+            return jnp.ones_like(x)
+
+    def zeros_like(self, x, cpu=True):
+        if cpu:
+            return np.zeros_like(x)
+        else:
+            return jnp.zeros_like(x)
+
+    def numel(self, v):
+        return np.size(v)
+
+    def asarray(self, v):
+        return jnp.asarray(v)
+
     def clean_redundancy(self):
         del self._train_loader
         del self._validation_loader