Home
last modified time | relevance | path

Searched refs:XLA_Shape (Results 1 – 12 of 12) sorted by relevance

/external/tensorflow/tensorflow/stream_executor/tpu/
Dc_api_decl.h219 typedef struct XLA_Shape { struct
223 XLA_Shape* tuple_shapes; // owned argument
226 } XLA_Shape; argument
230 XLA_Shape on_device_shape;
242 XLA_Shape shape;
246 XLA_Shape shape;
259 XLA_Shape dynamic_shape;
272 XLA_Shape* parameter_layouts;
273 XLA_Shape result_layout;
Dtpu_executor_c_api.h180 XLA_Shape* host_shape,
181 XLA_Shape* device_shape);
190 XLA_Shape* shape);
192 XLA_TransferManager* manager, XLA_Shape* host_shape, XLA_Shape* output,
202 SE_DeviceMemoryBase* elements, size_t elements_len, XLA_Shape* shape,
204 void TpuTransferManager_GetInfeedLayout(XLA_Shape* shape,
205 XLA_Shape* infeed_shape);
223 XLA_Shape* shape /*deprecated*/, XLA_Literal* c_literal, TF_Status* status);
305 XLA_Shape* c_shape);
338 XLA_Shape* serialized_xla_shape, int data_type, bool use_fast_memory,
[all …]
Dtpu_transfer_manager.cc54 XLA_Shape c_host_shape; in HostShapeToDeviceShape()
55 XLA_Shape c_device_shape; in HostShapeToDeviceShape()
133 XLA_Shape c_shape; in TransferLiteralFromOutfeed()
213 XLA_Shape c_shape; in GetByteSizeRequirement()
226 XLA_Shape c_host_shape; in ChooseCompactLayoutForShape()
228 XLA_Shape c_output; in ChooseCompactLayoutForShape()
279 XLA_Shape c_shape; in WriteSingleTupleIndexTable()
Dc_api_conversions.h46 xla::Shape FromC(const XLA_Shape* c_shape);
47 void ToC(const xla::Shape& xla_shape, XLA_Shape* c_shape);
48 void Free(XLA_Shape* c_shape);
Dtpu_executable.cc107 XLA_Shape c_host_shape; in HostShapeToDeviceShape()
108 XLA_Shape c_device_shape; in HostShapeToDeviceShape()
119 XLA_Shape c_shape; in ShapeSize()
Dc_api_conversions.cc212 void ToC(const xla::Shape& xla_shape, XLA_Shape* c_shape) { in ToC()
220 c_shape->tuple_shapes = new XLA_Shape[c_shape->ntuple_shapes]; in ToC()
233 xla::Shape FromC(const XLA_Shape* c_shape) { in FromC()
252 void Free(XLA_Shape* c_shape) { in Free()
441 new XLA_Shape[layout.parameter_count()]; in ToC()
/external/tensorflow/tensorflow/core/tpu/
Dtpu_ops_c_api.h178 XLA_Shape* host_shape, XLA_Shape* device_shape);
179 TFTPU_CAPI_EXPORT int64_t HardwareLayout_ShapeSize(XLA_Shape* shape);
180 TFTPU_CAPI_EXPORT int64_t HardwareLayout_ShapeSizeCompact(XLA_Shape* shape);
181 TFTPU_CAPI_EXPORT int64_t HardwareLayout_ShapeSizeCompactRaw(XLA_Shape* shape);
191 XLA_Shape* runtime_shape;
192 XLA_Shape* compile_time_shape;
Dtpu_execute.cc107 XLA_Shape c_host_shape; in HostShapeToDeviceShape()
108 XLA_Shape c_device_shape; in HostShapeToDeviceShape()
119 XLA_Shape c_shape; in ShapeSizeCompact()
128 XLA_Shape c_shape; in ShapeSizeCompactRaw()
239 XLA_Shape c_runtime_shape, c_compile_time_shape; in UpdateDynamicInputs()
Dtpu_on_demand_compiler.cc360 XLA_Shape c_shape; in ShapeSizeBytesFunction()
/external/tensorflow/tensorflow/core/tpu/kernels/xla/
Dinfeed_op.cc35 XLA_Shape c_shape; in GetTPUInfeedLayout()
36 XLA_Shape c_infeed_shape; in GetTPUInfeedLayout()
/external/tensorflow/tensorflow/compiler/jit/
Dxla_tpu_device.cc59 ApiConverter::StackHelper<XLA_Shape> se_shape(xla_shape); in TpuShapeRepresentation()
60 ApiConverter::StackHelper<XLA_Shape> tpu_shape; in TpuShapeRepresentation()
91 ApiConverter::StackHelper<XLA_Shape> se_shape(on_device_shape); in TpuPaddedShapeFn()
92 ApiConverter::StackHelper<XLA_Shape> tpu_shape; in TpuPaddedShapeFn()
/external/tensorflow/tensorflow/core/tpu/kernels/
Dinfeed_ops.cc58 XLA_Shape c_shape; in GetTPUInfeedLayout()
59 XLA_Shape c_infeed_shape; in GetTPUInfeedLayout()