(self, a)
| 1523 | return np.array(a) |
| 1524 | |
| 1525 | def _get_device(self, a): |
| 1526 | if self.jax_new_version: |
| 1527 | return list(a.devices())[0] |
| 1528 | else: |
| 1529 | return a.device_buffer.device() |
| 1530 | |
| 1531 | def _change_device(self, a, type_as): |
| 1532 | return jax.device_put(a, self._get_device(type_as)) |