Searched refs:shape_representation_fn_ (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | xla_device_context.cc | 90 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()
|
| D | xla_device.h | 73 return shape_representation_fn_; in shape_representation_fn()
83 XlaCompiler::ShapeRepresentationFn shape_representation_fn_; variable
245 const XlaCompiler::ShapeRepresentationFn shape_representation_fn_; variable
|
| D | xla_device_context.h | 83 return shape_representation_fn_; in shape_representation_fn()
114 XlaCompiler::ShapeRepresentationFn shape_representation_fn_; variable
|
| D | xla_device.cc | 153 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/ |
| D | mlir_hlo_to_hlo.cc | 452 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()
|