Home
last modified time | relevance | path

Searched refs:slice_dim (Results 1 – 11 of 11) sorted by relevance

/external/tensorflow/tensorflow/lite/micro/kernels/arc_mli/
Dmli_slicers.cc24 TensorSlicer::TensorSlicer(const mli_tensor* full_tensor, int slice_dim, in TensorSlicer() argument
31 sliceDim_(slice_dim),
41 if (i > slice_dim) {
43 } else if (i == slice_dim) {
57 if (i < slice_dim) {
59 } else if (i == slice_dim) {
65 sub_cfg_.sub_tensor_rank = full_tensor->rank - slice_dim;
Dmli_slicers.h26 TensorSlicer(const mli_tensor* full_tensor, int slice_dim, int slice_size,
/external/tensorflow/tensorflow/core/kernels/
Dscatter_nd_op.cc779 const int64 slice_dim = in ValidateUpdateShape() local
792 "Dimensions [", slice_dim, ",", params_shape.dims(), in ValidateUpdateShape()
794 "] must match dimensions [", slice_dim, ",", updates.dims(), in ValidateUpdateShape()
799 if (params_shape.dims() < slice_dim + (updates.dims() - batch_dim)) { in ValidateUpdateShape()
802 if (updates.dims() != batch_dim + params_shape.dims() - slice_dim) { in ValidateUpdateShape()
810 params_shape.dim_size(d + slice_dim)) { in ValidateUpdateShape()
820 int64* slice_dim, Index* num_updates, in PrepareAndValidateInputs() argument
862 *slice_dim = (indices_shape.dims() > 1) in PrepareAndValidateInputs()
872 for (int64 i = *slice_dim; i < total_nd; ++i) { in PrepareAndValidateInputs()
884 const int64 safe_slice_dim = (*slice_dim < 1) ? 1 : *slice_dim; in PrepareAndValidateInputs()
[all …]
/external/tensorflow/tensorflow/python/ops/
Dvariable_scope.py176 slice_dim = None
180 if slice_dim is not None:
184 "or {}.".format(shape, self.full_shape, i, slice_dim))
185 slice_dim = i
187 return slice_dim
740 slice_dim, num_slices = _get_slice_dim_and_num_slices(partitions)
758 _iter_slices(shape.as_list(), num_slices, slice_dim)):
2562 for slice_dim, num_slices in enumerate(slicing):
2567 slice_dim = 0
2569 return slice_dim, num_slices
[all …]
/external/tensorflow/tensorflow/python/eager/
Dpywrap_tensor.cc1016 PyObject* TFE_Py_TensorShapeSlice(PyObject* tensors, int slice_dim) { in TFE_Py_TensorShapeSlice() argument
1025 if (slice_dim < 0) { in TFE_Py_TensorShapeSlice()
1030 slice_dim) in TFE_Py_TensorShapeSlice()
1079 if (slice_dim >= num_dims) { in TFE_Py_TensorShapeSlice()
1082 tensorflow::strings::StrCat("Slice dimension (", slice_dim, in TFE_Py_TensorShapeSlice()
1089 int64_t dim = TFE_TensorHandleDim(handle, slice_dim, status.get()); in TFE_Py_TensorShapeSlice()
Dpywrap_tfe.h362 PyObject* TFE_Py_TensorShapeSlice(PyObject* tensors, int slice_dim);
/external/tensorflow/tensorflow/core/util/sparse/
Dsparse_tensor.h550 const int slice_dim = values_inserted_in_slice[slice_index]++; in Split() local
551 output_values_t[slice_index](slice_dim) = input_values_t(i); in Split()
554 output_indices_t[slice_index](slice_dim, j) = in Split()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Ddot_handler.cc942 auto slice_dim = in PartitionBaseCase() local
946 new_dims[slice_dim] /= 2; in PartitionBaseCase()
947 new_dims.insert(new_dims.begin() + slice_dim, 2); in PartitionBaseCase()
1008 auto slice_dim = in PartitionBaseCase() local
1012 slice_shape.set_dimensions(slice_dim, in PartitionBaseCase()
1013 dot->shape().dimensions(slice_dim) / 2); in PartitionBaseCase()
1016 cw_start_indices[slice_dim] = dot->shape().dimensions(slice_dim) / 2; in PartitionBaseCase()
Dspmd_partitioner_util.h278 int64 slice_dim, int64 k);
Dspmd_partitioner_util.cc1262 int64 slice_dim, int64 k) { in SliceFirstK() argument
1268 limit_indices[slice_dim] = k; in SliceFirstK()
1270 output_shape.set_dimensions(slice_dim, k); in SliceFirstK()
/external/tensorflow/tensorflow/python/
Dtfe_wrapper.cc1031 [](const py::handle& tensors, int slice_dim) { in PYBIND11_MODULE() argument
1033 TFE_Py_TensorShapeSlice(tensors.ptr(), slice_dim)); in PYBIND11_MODULE()
1036 int slice_dim) { in PYBIND11_MODULE() argument