Searched refs:make_device_array (Results 1 – 3 of 3) sorted by relevance
| /dports/math/py-jax/jax-0.2.9/jax/_src/ |
| H A D | dlpack.py | 69 …return xla.make_device_array(aval, buf.device(), lazy.array(aval.shape), buf) # pytype: disable=a…
|
| /dports/math/py-jax/jax-0.2.9/jax/interpreters/ |
| H A D | xla.py | 117 return partial(make_device_array, raise_to_shaped(aval), device,
1028 def make_device_array( function
1316 return make_device_array(x.aval, device, x._lazy_expr,
1346 return make_device_array(x.aval, x._device, lazy.array(x.aval.shape), result)
|
| H A D | pxla.py | 593 return xla.make_device_array(aval, None, lazy.array(aval.shape), buf)
|