Home
last modified time | relevance | path

Searched refs:_DeviceArray (Results 1 – 4 of 4) sorted by relevance

/dports/math/py-jax/jax-0.2.9/jax/interpreters/
H A Dxla.py1046 return _DeviceArray(aval, device, lazy_expr, device_buffer)
1055 return type_x is _DeviceArray or type_x is _CppDeviceArray
1058 class _DeviceArray(DeviceArray): # type: ignore class
1169 for device_array in [_DeviceArray, _CppDeviceArray]:
1289 for device_array in [_CppDeviceArray, _DeviceArray]:
1301 xb.register_constant_handler(_DeviceArray, _device_array_constant_handler)
1308 device_put_handlers[_DeviceArray] = _device_put_device_array
1310 …ce(x: Union[DeviceArrayProtocol, _DeviceArray], device: Optional[xc.Device]) -> Union[DeviceArrayP… argument
1333 return _DeviceArray(x.aval, device, x._lazy_expr, moved_buf)
1351 ) -> Callable[[_DeviceArray], PyLocalBuffer]:
[all …]
H A Dpxla.py389 shard_arg_handlers[xla._DeviceArray] = _shard_device_array
475 class ShardedDeviceArray(xla._DeviceArray):
/dports/math/py-jax/jax-0.2.9/jax/_src/lax/
H A Dlax.py1480 return xla._DeviceArray(aval, None, lazy_expr, xla.DeviceConstant())
1502 return xla._DeviceArray(aval, None, lazy_expr, xla.DeviceConstant())
1519 return xla._DeviceArray(aval, None, lazy_expr, xla.DeviceConstant())
1533 return xla._DeviceArray(aval, None, lazy_expr, xla.DeviceConstant())
1958 [xla._CppDeviceArray, xla._DeviceArray, pxla.ShardedDeviceArray]):
1960 ad_util.jaxval_zeros_likers[xla._DeviceArray] = zeros_like_array
3276 return xla._DeviceArray(aval, operand._device, lazy_expr, operand.device_buffer)
3583 return xla._DeviceArray(aval, operand._device, lazy_expr, operand.device_buffer)
3700 return xla._DeviceArray(aval, operand._device, lazy_expr, operand.device_buffer)
/dports/math/py-jax/jax-0.2.9/jax/_src/numpy/
H A Dlax_numpy.py47 from jax.interpreters.xla import DeviceArray, _DeviceArray, _CppDeviceArray
5205 for device_array in [_DeviceArray, _CppDeviceArray]:
5232 setattr(_DeviceArray, "__array_module__", __array_module__)
5240 for device_array in [_DeviceArray, _CppDeviceArray]:
5249 setattr(_DeviceArray, "compress", _compress_method)
5269 setattr(_DeviceArray, "_multi_slice", _multi_slice)
5391 setattr(_DeviceArray, "at", property(_IndexUpdateHelper))