(self, a, type_as=None)
| 1532 | return jax.device_put(a, self._get_device(type_as)) |
| 1533 | |
| 1534 | def _from_numpy(self, a, type_as=None): |
| 1535 | if isinstance(a, float): |
| 1536 | a = np.array(a) |
| 1537 | if type_as is None: |
| 1538 | return jnp.array(a) |
| 1539 | else: |
| 1540 | return self._change_device(jnp.array(a).astype(type_as.dtype), type_as) |
| 1541 | |
| 1542 | def set_gradients(self, val, inputs, grads): |
| 1543 | from jax.flatten_util import ravel_pytree |
nothing calls this directly
no test coverage detected