Searched refs:HostShapeToDeviceShape (Results 1 – 16 of 16) sorted by relevance
34 client->backend().transfer_manager()->HostShapeToDeviceShape(shape), in MakeArray()68 client->backend().transfer_manager()->HostShapeToDeviceShape(a_shape)); in TEST()70 client->backend().transfer_manager()->HostShapeToDeviceShape(b_shape)); in TEST()72 client->backend().transfer_manager()->HostShapeToDeviceShape(c_shape)); in TEST()110 client->backend().transfer_manager()->HostShapeToDeviceShape( in TEST()
161 const Shape on_device_shape = HostShapeToDeviceShape(literal.shape()); in TransferArrayToDeviceAsync()183 if (!Shape::Equal().MinorToMajorOnlyInLayout()(HostShapeToDeviceShape(shape), in TransferArrayFromDevice()187 ShapeUtil::HumanString(HostShapeToDeviceShape(shape))); in TransferArrayFromDevice()395 Shape on_device_shape = HostShapeToDeviceShape(on_host_shape); in AllocateScopedShapedBuffer()
53 virtual Shape HostShapeToDeviceShape(const Shape& host_shape) const { in HostShapeToDeviceShape() function
59 Shape HostShapeToDeviceShape(const Shape& host_shape) override;
82 virtual Shape HostShapeToDeviceShape(const Shape& host_shape) = 0;
41 xla::Shape HostShapeToDeviceShape(
106 Shape TpuExecutable::HostShapeToDeviceShape(const Shape& host_shape) { in HostShapeToDeviceShape() function in xla::TpuExecutable
63 Shape device_shape = HostShapeToDeviceShape(host_shape); in AllocateOutputMemoryWithInputReuse()
52 xla::Shape TpuTransferManager::HostShapeToDeviceShape( in HostShapeToDeviceShape() function in tensorflow::tpu::TpuTransferManager
106 xla::Shape HostShapeToDeviceShape(const xla::Shape& host_shape) { in HostShapeToDeviceShape() function208 HostShapeToDeviceShape(compile_time_shapes[i]); in UpdateDynamicInputs()467 backend->transfer_manager()->HostShapeToDeviceShape( in TPUExecute()
198 Shape HostShapeToDeviceShape(const Shape& host_shape) override { in HostShapeToDeviceShape() function in xla::__anoncef202870111::TpuExecutable
47 client->backend().transfer_manager()->HostShapeToDeviceShape( in AllocateShapedBuffer()
237 transfer_manager->HostShapeToDeviceShape(shape); in PopulateInputs()
130 backend->transfer_manager()->HostShapeToDeviceShape( in GetExecutableAliasConfig()
87 xla::Shape on_device_shape = transfer_manager->HostShapeToDeviceShape(shape); in AllocateScopedShapedBuffer()
282 transfer_manager->HostShapeToDeviceShape(input_host_shape)); in BuildComputationInputs()