Home
last modified time | relevance | path

Searched refs:tensor_dim_index (Results 1 – 3 of 3) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/
H A Dsave_restore_util.cc172 for (int tensor_dim_index = 0; tensor_dim_index < sharding_spec_strs.size(); in SliceSpecOnDevice() local
173 ++tensor_dim_index) { in SliceSpecOnDevice()
174 const std::string& mesh_dim = sharding_spec_strs[tensor_dim_index]; in SliceSpecOnDevice()
178 int64_t per_slice_size = global_shape[tensor_dim_index] / dim_size; in SliceSpecOnDevice()
180 slice_specs[tensor_dim_index] = absl::StrCat(start, ",", per_slice_size); in SliceSpecOnDevice()
/aosp_15_r20/external/tensorflow/tensorflow/lite/tools/benchmark/
H A Dbenchmark_test.cc116 int tensor_dim_index, in CheckInputTensorValue() argument
118 StringRef tensor_value = GetString(input_tensor, tensor_dim_index); in CheckInputTensorValue()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dlayout_propagation_v2.cc487 const int tensor_dim_index = *it.second.begin(); in GetMostShardedLayout() local
488 dim_to_layout_map[tensor_dim_index].insert(it.first); in GetMostShardedLayout()