| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | utils.cc | 150 absl::optional<std::vector<Shape>>& argument_layouts, in DetermineArgumentLayoutsFromCompileOptions() argument
155 if (!argument_layouts) { in DetermineArgumentLayoutsFromCompileOptions()
156 argument_layouts.emplace(program_shape.parameters()); in DetermineArgumentLayoutsFromCompileOptions()
157 for (Shape& shape : *argument_layouts) { in DetermineArgumentLayoutsFromCompileOptions()
160 } else if (argument_layouts->size() != program_shape.parameters_size()) { in DetermineArgumentLayoutsFromCompileOptions()
164 argument_layouts->size(), program_shape.parameters_size()); in DetermineArgumentLayoutsFromCompileOptions()
166 argument_layout_pointers->reserve(argument_layouts->size()); in DetermineArgumentLayoutsFromCompileOptions()
190 CHECK_EQ(sharded_shapes.first.size(), argument_layouts->size()); in DetermineArgumentLayoutsFromCompileOptions()
191 for (int i = 0; i < argument_layouts->size(); ++i) { in DetermineArgumentLayoutsFromCompileOptions()
192 Shape* layout = &(*argument_layouts)[i]; in DetermineArgumentLayoutsFromCompileOptions()
|
| D | utils.h | 44 absl::optional<std::vector<Shape>>& argument_layouts,
|
| D | pjrt_client.h | 117 absl::optional<std::vector<Shape>> argument_layouts; member
|
| D | pjrt_stream_executor_client.cc | 2145 options.argument_layouts, &options.executable_build_options, in Compile()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | local_service.cc | 103 const absl::Span<const Shape* const> argument_layouts, in CompileExecutables() argument
110 if (argument_layouts.size() != program_shape.parameters_size()) { in CompileExecutables()
113 program_shape.parameters_size(), argument_layouts.size()); in CompileExecutables()
116 for (int i = 0; i < argument_layouts.size(); ++i) { in CompileExecutables()
117 const Shape& argument_shape = *argument_layouts[i]; in CompileExecutables()
151 CreateModuleConfig(program_shape, argument_layouts, &execution_options)); in CompileExecutables()
|
| D | local_service.h | 52 const absl::Span<const Shape* const> argument_layouts,
|
| D | compile_only_service.h | 44 std::vector<const Shape*> argument_layouts; member
|
| D | compile_only_service.cc | 107 instance.argument_layouts, &execution_options, &options)); in CompileAheadOfTime()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | compile_only_client.cc | 47 service_instance.argument_layouts = instance.argument_layouts; in CompileAheadOfTime()
|
| D | compile_only_client.h | 44 std::vector<const Shape*> argument_layouts; member
|
| D | local_client.h | 145 const absl::Span<const Shape* const> argument_layouts,
|
| D | local_client.cc | 358 const absl::Span<const Shape* const> argument_layouts, in Compile() argument
385 computation, argument_layouts, updated_options)); in Compile()
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | local_client_test_base.cc | 192 std::vector<const Shape*> argument_layouts(arguments.size()); in ExecuteLocally() local
194 argument_layouts[i] = &arguments[i]->on_host_shape(); in ExecuteLocally()
198 local_client_->Compile(computation, argument_layouts, build_options)); in ExecuteLocally()
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | xrt_compile_ops.cc | 152 std::vector<xla::Shape> argument_layouts( in Compile() local
157 argument_layouts[i] = xla::Shape(config.program_shape().parameters(i)); in Compile()
158 argument_layout_ptrs[i] = &argument_layouts[i]; in Compile()
|
| /external/tensorflow/tensorflow/compiler/xla/tools/ |
| D | replay_computation.cc | 117 std::vector<Shape> argument_layouts; in CompileExecutable() local
118 argument_layouts.reserve( in CompileExecutable()
123 argument_layouts.push_back(Shape(param)); in CompileExecutable()
124 argument_layout_ptrs.push_back(&argument_layouts.back()); in CompileExecutable()
|
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | xla_compilation_cache.cc | 163 std::vector<const xla::Shape*> argument_layouts( in BuildExecutable() local
166 argument_layouts[i] = &result.xla_input_shapes[i]; in BuildExecutable()
179 client_->Compile(*result.computation, argument_layouts, build_options)); in BuildExecutable()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_util.cc | 59 instance.argument_layouts.push_back(&shape); in BuildAotXlaComputationInstance()
|
| /external/tensorflow/tensorflow/compiler/aot/ |
| D | compile.cc | 87 instance.argument_layouts = std::move(arg_layout_ptrs); in CompileXla()
|
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| D | tpu_client_extension.cc | 136 computation, options.argument_layouts, in PYBIND11_MODULE()
|
| D | tpu_client.h | 293 absl::optional<std::vector<Shape>> argument_layouts,
|
| D | tpu_client.cc | 796 absl::optional<std::vector<Shape>> argument_layouts, in Compile() argument
834 CHECK(!argument_layouts.has_value()); in Compile()
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | xla_compiler.cc | 449 .def_readwrite("argument_layouts", &CompileOptions::argument_layouts) in BuildXlaCompilerSubmodule()
|
| D | xla_client_test.py | 1873 options.argument_layouts = [
|