Home
last modified time | relevance | path

Searched refs:type_is_device_array (Results 1 – 3 of 3) sorted by relevance

/dports/math/py-jax/jax-0.2.9/jax/interpreters/
H A Dxla.py859 return [_copy_device_array_to_device(x, device) if type_is_device_array(x)
1049 def type_is_device_array(x): function
1287 return type_is_device_array(x) and type(x.device_buffer) is DeviceConstant
1382 if type_is_device_array(x):
/dports/math/py-jax/jax-0.2.9/jax/
H A Dapi.py312 all(xla.type_is_device_array(x) for x in out_flat) and
/dports/math/py-jax/jax-0.2.9/jax/_src/lax/
H A Dlax.py3267 if xla.type_is_device_array(operand) and np.all(
3575 if xla.type_is_device_array(operand) and dimensions is None:
3694 if xla.type_is_device_array(operand):