Searched defs:host_shape (Results 1 – 8 of 8) sorted by relevance
67 ExecutionInput(xla::Shape shape, xla::Shape host_shape) in ExecutionInput()78 xla::Shape host_shape) in ExecutionInput()93 const Shape& host_shape() const { in host_shape() function132 void SetHostShape(xla::Shape host_shape) { in SetHostShape()
53 virtual Shape HostShapeToDeviceShape(const Shape& host_shape) const { in HostShapeToDeviceShape()
106 Shape TpuExecutable::HostShapeToDeviceShape(const Shape& host_shape) { in HostShapeToDeviceShape()
55 const Shape& host_shape, const HloInputOutputAliasConfig& alias_config, in AllocateOutputMemoryWithInputReuse()
415 se::DeviceMemoryAllocator* allocator, xla::Shape* host_shape, in ExpandTreeOfTuples()462 xla::Shape host_shape; in MakeTuple() local
321 const xla::Shape& host_shape = in BuildComputationInputs() local663 xla::Shape host_shape(executable.input_shapes(0)); in DoWork() local
198 Shape HostShapeToDeviceShape(const Shape& host_shape) override { in HostShapeToDeviceShape()
106 xla::Shape HostShapeToDeviceShape(const xla::Shape& host_shape) { in HostShapeToDeviceShape()