Searched refs:_DeviceArray (Results 1 – 4 of 4) sorted by relevance
/dports/math/py-jax/jax-0.2.9/jax/interpreters/ |
H A D | xla.py | 1046 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 D | pxla.py | 389 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 D | lax.py | 1480 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 D | lax_numpy.py | 47 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))
|