| /external/tensorflow/tensorflow/dtensor/python/ |
| D | __init__.py | 17 from tensorflow.dtensor.python import gen_dtensor_ops as ops
18 from tensorflow.dtensor.python import mesh_util
19 from tensorflow.dtensor.python import tpu_util
21 from tensorflow.dtensor.python.api import call_with_layout
22 from tensorflow.dtensor.python.api import check_layout
23 from tensorflow.dtensor.python.api import client_id
24 from tensorflow.dtensor.python.api import copy_to_mesh
25 from tensorflow.dtensor.python.api import device_name
26 from tensorflow.dtensor.python.api import fetch_layout
27 from tensorflow.dtensor.python.api import full_job_name
[all …]
|
| D | BUILD | 7 "//tensorflow/dtensor:dtensor-internal",
23 "//tensorflow/dtensor:dtensor-users",
61 "//tensorflow/dtensor/cc:dtensor_ops",
62 "//tensorflow/dtensor/cc:dtensor_tpu_ops",
71 "//tensorflow/dtensor/proto:layout_proto_py_pb2",
110 "//tensorflow/dtensor/proto:layout_proto_py_pb2",
179 "//tensorflow/dtensor:dtensor-users",
201 "//tensorflow/dtensor:dtensor-users",
|
| D | dtensor_device.py | 26 from tensorflow.dtensor.python import gen_dtensor_ops
27 from tensorflow.dtensor.python import layout as layout_lib
243 def unpack(self, dtensor: Any) -> Sequence[Any]:
266 if issubclass(type(dtensor), resource_variable_ops.BaseResourceVariable):
272 dtensor,
279 dtensor,
292 def fetch_layout(self, dtensor: Any) -> layout_lib.Layout:
306 if issubclass(type(dtensor), resource_variable_ops.BaseResourceVariable):
307 dtensor = dtensor.read_value()
311 dtensor,
|
| D | save_restore.py | 20 from tensorflow.dtensor.python import api
21 from tensorflow.dtensor.python import d_variable
22 from tensorflow.dtensor.python import gen_dtensor_ops
23 from tensorflow.dtensor.python import layout as layout_lib
24 from tensorflow.dtensor.python import mesh_util
|
| /external/tensorflow/tensorflow/dtensor/mlir/utils/ |
| D | BUILD | 7 "//tensorflow/dtensor:dtensor-internal",
34 "//tensorflow/dtensor/cc:constants",
35 "//tensorflow/dtensor/cc:dstatus",
36 "//tensorflow/dtensor/cc:dtensor_utils",
37 "//tensorflow/dtensor/cc:tensor_layout",
38 "//tensorflow/dtensor/mlir:collectives_common",
39 "//tensorflow/dtensor/mlir:create_dtensor_mlir_passes",
40 "//tensorflow/dtensor/mlir:device_utils",
41 "//tensorflow/dtensor/mlir:dtensor_location",
42 "//tensorflow/dtensor/mlir:group_assignment",
[all …]
|
| /external/tensorflow/tensorflow/dtensor/mlir/ |
| D | Passes.td | 22 : Pass<"dtensor-op-to-device-cluster", "mlir::func::FuncOp"> {
30 : Pass<"dtensor-device-mesh-cluster-coarsening", "mlir::func::FuncOp"> {
38 : Pass<"dtensor-constant-folding", "mlir::func::FuncOp"> {
46 : Pass<"dtensor-allreduce-sum-optimization", "mlir::func::FuncOp"> {
54 : Pass<"dtensor-allreduce-scatter-optimization", "mlir::func::FuncOp"> {
62 : Pass<"dtensor-allreduce-combine-optimization", "mlir::func::FuncOp"> {
70 : Pass<"dtensor-mixed-precision-reduce", "mlir::func::FuncOp"> {
78 : Pass<"dtensor-dce", "mlir::func::FuncOp"> {
86 : Pass<"dtensor-undo-merge-const-across-mesh", "mlir::func::FuncOp"> {
94 : Pass<"dtensor-set-default-sharding", "mlir::func::FuncOp"> {
[all …]
|
| D | BUILD | 10 "//tensorflow/dtensor:dtensor-internal",
72 "//tensorflow/dtensor/mlir/dtensor_dialect:ir/dtensor_attributes",
102 "//tensorflow/dtensor/cc:dstatus",
103 "//tensorflow/dtensor/cc:tensor_layout",
116 "//tensorflow/dtensor/cc:tensor_layout",
127 "//tensorflow/dtensor/cc:dstatus",
128 "//tensorflow/dtensor/cc:tensor_layout",
244 "//tensorflow/dtensor/cc:constants",
245 "//tensorflow/dtensor/cc:dtensor_utils",
246 "//tensorflow/dtensor/cc:tensor_layout",
[all …]
|
| D | dtensor_location_test.cc | 38 loc = tensorflow::dtensor::DTensorLocation(loc, "test.cc", 21); in TEST()
47 EXPECT_EQ(tensorflow::dtensor::DTensorLocationToString(loc), stack); in TEST()
53 test_loc = tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 21); in TEST()
54 test_loc = tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 22); in TEST()
55 test_loc = tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 23); in TEST()
56 test_loc = tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 24); in TEST()
73 EXPECT_EQ(tensorflow::dtensor::DTensorLocationToString(test_loc), stack); in TEST()
81 test_loc = tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 21); in TEST()
97 EXPECT_EQ(tensorflow::dtensor::DTensorLocationToString(test_loc), stack); in TEST()
106 tensorflow::dtensor::DTensorLocation(test_loc, "test.cc", 21, "nested"); in TEST()
[all …]
|
| D | collectives.h | 28 namespace dtensor {
35 const dtensor::Layout& src_layout, const dtensor::Layout& tgt_layout,
52 mlir::Value input, const dtensor::Layout& src_layout,
53 const dtensor::Layout& tgt_layout,
58 mlir::OpBuilder& builder, const dtensor::Layout& output_layout,
|
| D | dtensor_location.h | 30 namespace dtensor {
55 ::tensorflow::dtensor::DTensorLocation(loc, __FILE__, __LINE__)
59 ::tensorflow::dtensor::DTensorLocation(loc, __FILE__, __LINE__, name)
|
| /external/eigen/unsupported/test/ |
| D | cxx11_tensor_casts.cpp | 45 Tensor<double, 2> dtensor(20,30); in test_vectorized_cast() local
46 dtensor.setRandom(); in test_vectorized_cast()
49 dtensor = itensor.cast<double>(); in test_vectorized_cast()
54 VERIFY_IS_EQUAL(dtensor(i,j), static_cast<double>(ftensor(i,j))); in test_vectorized_cast()
64 Tensor<double, 2> dtensor(20,30); in test_float_to_int_cast() local
65 dtensor = dtensor.random() * 1000.0; in test_float_to_int_cast()
68 Tensor<int, 2> i2tensor = dtensor.cast<int>(); in test_float_to_int_cast()
73 VERIFY_IS_EQUAL(i2tensor(i,j), static_cast<int>(dtensor(i,j))); in test_float_to_int_cast()
81 Tensor<double, 2> dtensor(20, 30); in test_big_to_small_type_cast() local
82 dtensor.setRandom(); in test_big_to_small_type_cast()
[all …]
|
| /external/tensorflow/tensorflow/python/ |
| D | pywrap_dtensor_device.cc | 31 using tensorflow::dtensor::AddMesh;
32 using tensorflow::dtensor::AllocateDTensorDevice;
33 using tensorflow::dtensor::ClearTPUCoreIDs;
34 using tensorflow::dtensor::ExperimentalClearDefaultLayout;
35 using tensorflow::dtensor::ExperimentalClearDefaultMesh;
36 using tensorflow::dtensor::ExperimentalSetDefaultLayout;
37 using tensorflow::dtensor::ExperimentalSetDefaultMesh;
38 using tensorflow::dtensor::FetchLayout;
39 using tensorflow::dtensor::GetFunctionCacheHitAndMissCount;
40 using tensorflow::dtensor::IsSparseDTensor;
[all …]
|
| /external/tensorflow/tensorflow/dtensor/mlir/dtensor_dialect/ir/ |
| D | dtensor_attributes.cc | 25 namespace dtensor { namespace
30 using Mesh = tensorflow::dtensor::Mesh;
59 using Layout = tensorflow::dtensor::Layout;
82 tensorflow::dtensor::Layout layout) { in get()
91 addAttributes<dtensor::MeshAttr, dtensor::LayoutAttr>(); in registerAttributes()
|
| D | dtensor_attributes.h | 26 namespace dtensor {
38 using Mesh = tensorflow::dtensor::Mesh;
52 using Layout = tensorflow::dtensor::Layout;
53 using Mesh = tensorflow::dtensor::Mesh;
|
| D | ops.cc | 33 namespace dtensor { namespace
70 using Mesh = tensorflow::dtensor::Mesh; in ParseMeshAttr()
71 using MeshOr = tensorflow::dtensor::StatusOr<Mesh>; in ParseMeshAttr()
102 using Layout = tensorflow::dtensor::Layout; in ParseLayoutAttr()
103 using LayoutOr = tensorflow::dtensor::StatusOr<Layout>; in ParseLayoutAttr()
|
| /external/tensorflow/tensorflow/dtensor/cc/ |
| D | dtensor_graph_to_mlir_pass.cc | 52 logging_enabled_ = dtensor::MaybeEnableLogging(&pass_manager_); in DTensorMlirPassRunner()
64 dtensor::CreateDTensorMLIRPass(pipeline_options, &pass_manager_); in DTensorMlirPassRunner()
102 if (!is_func && !dtensor::LogOpByOp()) in RunOnGraph()
103 module->setAttr(dtensor::kDoNotLog, mlir::UnitAttr::get(&context_)); in RunOnGraph()
109 dtensor::kCacheKey, in RunOnGraph()
116 if (logging_enabled_ && !module->hasAttr(dtensor::kDoNotLog)) in RunOnGraph()
|
| D | dstatus.h | 31 namespace dtensor {
58 ::tensorflow::dtensor::WithContext(dstatus, __FILE__, __LINE__, #__VA_ARGS__);
65 return ::tensorflow::dtensor::WithContext(_status, __FILE__, __LINE__); \
74 return ::tensorflow::dtensor::WithContext(_status, __FILE__, __LINE__, \
85 return ::tensorflow::dtensor::WithContext(statusor.status(), __FILE__, \
|
| D | BUILD | 11 "//tensorflow/dtensor:dtensor-internal",
55 "//tensorflow/dtensor/proto:layout_proto_cc",
82 "//tensorflow/dtensor/proto:layout_proto_cc",
245 "//tensorflow/dtensor/mlir:dtensor_mlir_passes",
246 "//tensorflow/dtensor/mlir:tf_dtensor_dialect",
247 "//tensorflow/dtensor/mlir/dtensor_dialect:Dialect",
295 "//tensorflow/dtensor/mlir:layout_parsing",
296 "//tensorflow/dtensor/proto:layout_proto_cc",
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | list_ops.py | 232 def _TensorListStackGrad(unused_op, dtensor): argument
233 return tensor_list_from_tensor(dtensor, element_shape=dtensor.shape[1:]), None
238 def _TensorListConcatGrad(op, dtensor, unused_dlengths): argument
241 dtensor,
324 def _TensorListGatherGrad(op, dtensor): argument
330 dlist = tensor_list_reserve(element_shape, num_elements, dtensor.dtype)
332 tensor=dtensor, indices=indices, input_handle=dlist)
342 dtensor = gen_list_ops.tensor_list_gather(
348 return dtensor, None, None, None
350 return dtensor, None, None
[all …]
|
| /external/tensorflow/tensorflow/dtensor/ |
| D | BUILD | 9 name = "dtensor-users",
13 "//third_party/py/keras/dtensor/...",
18 name = "dtensor-internal",
23 "//tensorflow/dtensor/...",
|
| /external/tensorflow/tensorflow/dtensor/proto/ |
| D | BUILD | 5 "//tensorflow/dtensor:dtensor-internal",
6 "//tensorflow/dtensor:dtensor-users",
|
| /external/tensorflow/tensorflow/dtensor/mlir/dtensor_dialect/ |
| D | BUILD | 8 "//tensorflow/dtensor:dtensor-internal",
68 "//tensorflow/dtensor/cc:tensor_layout",
84 "//tensorflow/dtensor/cc:tensor_layout",
85 "//tensorflow/dtensor/proto:layout_proto_cc",
|
| /external/tensorflow/tensorflow/tools/def_file_filter/ |
| D | symbols_pybind.txt | 514 [//tensorflow/dtensor/cc:dtensor_device_cc] # DTensor
515 tensorflow::dtensor::AllocateDTensorDevice
516 tensorflow::dtensor::AddMesh
517 tensorflow::dtensor::ExperimentalSetDefaultLayout
518 tensorflow::dtensor::ExperimentalClearDefaultLayout
519 tensorflow::dtensor::ExperimentalSetDefaultMesh
520 tensorflow::dtensor::ExperimentalClearDefaultMesh
521 tensorflow::dtensor::SetSameShapePolicy
522 tensorflow::dtensor::SetTPUCoreIDs
523 tensorflow::dtensor::ClearTPUCoreIDs
[all …]
|
| /external/tensorflow/tensorflow/dtensor/mlir/ir/ |
| D | tf_dtensor.cc | 89 const tensorflow::dtensor::Layout input_layout = op.input_layout(); in verify()
90 const tensorflow::dtensor::Layout output_layout = op.output_layout(); in verify()
99 tensorflow::dtensor::Layout::IsShardedDimension( in verify()
145 const tensorflow::dtensor::Layout input_layout = op.input_layout(); in verify()
146 const tensorflow::dtensor::Layout output_layout = op.output_layout(); in verify()
155 tensorflow::dtensor::Layout::IsShardedDimension( in verify()
|
| /external/tensorflow/tensorflow/python/tools/api/generator/ |
| D | api_gen.bzl | 48 "tensorflow.dtensor.python.api",
49 "tensorflow.dtensor.python.d_checkpoint",
50 "tensorflow.dtensor.python.d_variable",
51 "tensorflow.dtensor.python.input_util",
52 "tensorflow.dtensor.python.layout",
53 "tensorflow.dtensor.python.mesh_util",
54 "tensorflow.dtensor.python.save_restore",
55 "tensorflow.dtensor.python.tpu_util",
|