Searched refs:DevicePutResult (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | jax_jit.cc | 567 using DevicePutFunc = std::function<xla::StatusOr<DevicePutResult>(
570 DevicePutResult HandleBool(py::handle h, xla::PjRtDevice* to_device, in HandleBool()
572 return DevicePutResult(ConvertToScalarBuffer<bool, py::bool_>( in HandleBool()
577 DevicePutResult HandleInt(py::handle obj, xla::PjRtDevice* to_device, in HandleInt()
580 return DevicePutResult(ConvertToScalarBuffer<xla::int64, py::int_>( in HandleInt()
584 return DevicePutResult(ConvertToScalarBuffer<int, py::int_>( in HandleInt()
591 xla::StatusOr<DevicePutResult> HandleFloat(py::handle h, in HandleFloat()
596 return DevicePutResult(ConvertToScalarBuffer<double, py::float_>( in HandleFloat()
600 return DevicePutResult(ConvertToScalarBuffer<float, py::float_>( in HandleFloat()
607 xla::StatusOr<DevicePutResult> HandleComplex(py::handle h, in HandleComplex()
[all …]
|
| D | jax_jit.h | 141 struct DevicePutResult { struct
142 explicit DevicePutResult(xla::PjRtBuffer* b, bool weak_type) in DevicePutResult() function
144 DevicePutResult(std::unique_ptr<xla::PjRtBuffer> new_buffer, bool weak_type) in DevicePutResult() argument
168 xla::StatusOr<DevicePutResult> DevicePut(pybind11::handle arg, argument
|