Home
last modified time | relevance | path

Searched refs:shape_representation_fn_ (Results 1 – 5 of 5) sorted by relevance

/external/tensorflow/tensorflow/compiler/jit/
Dxla_device_context.cc90 shape_representation_fn_(std::move(shape_representation_fn)), in XlaDeviceContext()
95 if (!shape_representation_fn_) { in XlaDeviceContext()
96 shape_representation_fn_ = in XlaDeviceContext()
139 shape_representation_fn_(device_tensor->shape(), device_tensor->dtype(), in CopyCPUTensorToDevice()
Dxla_device.h73 return shape_representation_fn_; in shape_representation_fn()
83 XlaCompiler::ShapeRepresentationFn shape_representation_fn_; variable
245 const XlaCompiler::ShapeRepresentationFn shape_representation_fn_; variable
Dxla_device_context.h83 return shape_representation_fn_; in shape_representation_fn()
114 XlaCompiler::ShapeRepresentationFn shape_representation_fn_; variable
Dxla_device.cc153 shape_representation_fn_(std::move(shape_representation_fn)), in Metadata()
213 shape_representation_fn_(options.shape_representation_fn), in XlaDevice()
345 device_to_device_streams, client, shape_representation_fn_, in GetDeviceContextLocked()
353 client, shape_representation_fn_, thread_pool_.get(), true); in GetDeviceContextLocked()
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dmlir_hlo_to_hlo.cc452 shape_representation_fn_(shape_representation_fn), in ConvertToHloModule()
454 if (!shape_representation_fn_) in ConvertToHloModule()
455 shape_representation_fn_ = tensorflow::IdentityShapeRepresentationFn(); in ConvertToHloModule()
538 tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn_; member in mlir::__anonf0ba238b0211::ConvertToHloModule
1322 builder, returns[index], return_shape, shape_representation_fn_, in Lower()
1462 auto arg_shape_status = shape_representation_fn_(arg_tensor_shape, dtype, in SetEntryTupleShapesAndLeafReplication()
1494 shape_representation_fn_, &(*arg_shapes)[arg_sharding.index()]); in SetEntryTupleShardings()