HomeSort by relevance Sort by last modified time
    Searched refs:suffix_dim_size (Results 1 - 5 of 5) sorted by null

  /external/tensorflow/tensorflow/core/kernels/
split_lib_gpu.cu.cc 73 int32 split_dim_size, int32 suffix_dim_size,
82 int32 size = prefix_dim_size * split_dim_size * suffix_dim_size;
87 int32 i = offset / (split_dim_size * suffix_dim_size);
88 int32 j = (offset % (split_dim_size * suffix_dim_size)) / suffix_dim_size;
89 int32 k = offset % suffix_dim_size;
94 // [prefix_dim_size][piece_size][suffix_dim_size].
98 int32 output_offset = i * piece_size * suffix_dim_size +
99 (j % piece_size) * suffix_dim_size + k;
168 const T* input, int32 prefix_dim_size, int32 suffix_dim_size,
    [all...]
split_op.cc 111 // Caller must ensure that dim_size and suffix_dim_size are <
116 IndexType suffix_dim_size = 1;
118 suffix_dim_size *= static_cast<IndexType>(input_shape.dim_size(i));
120 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size);
153 Eigen::DenseIndex suffix_dim_size;
155 std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
158 input.shaped<T, 3>({prefix_dim_size, split_dim_size, suffix_dim_size});
166 prefix_dim_size, split_dim_output_size, suffix_dim_size};
178 split_dim_output_size, suffix_dim_size, &sizes,
185 if (prefix_dim_size * split_dim_output_size * suffix_dim_size > 0)
    [all...]
split_v_op.cc 165 // Caller must ensure that dim_size and suffix_dim_size are <
170 IndexType suffix_dim_size = 1; local
172 suffix_dim_size *= static_cast<IndexType>(input_shape.dim_size(i));
174 return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size);
208 Eigen::DenseIndex suffix_dim_size;
210 std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
213 input.shaped<T, 3>({prefix_dim_size, split_dim_size, suffix_dim_size});
237 suffix_dim_size, use_parallelism_between_outputs,
247 prefix_dim_size, split_sizes_vec[i], suffix_dim_size};
251 {prefix_dim_size, split_sizes_vec[i], suffix_dim_size});
    [all...]
one_hot_op.cc 101 // suffix_dim_size == # of elements after the axis
106 TI suffix_dim_size = indices_shape.num_elements() / prefix_dim_size; variable
108 // Split indices into matrix of size prefix_dim_size x suffix_dim_size
110 indices.shaped<TI, 2>({prefix_dim_size, suffix_dim_size});
112 // prefix_dim_size x depth x suffix_dim_size.
114 output->shaped<T, 3>({prefix_dim_size, depth_v, suffix_dim_size});
batch_kernels.cc 144 int64 suffix_dim_size = 1; local
146 suffix_dim_size *= input.shape().dim_size(i);
149 input.shaped<T, 3>({1, input.shape().dim_size(0), suffix_dim_size});
158 auto output_shaped = output.shaped<T, 3>({1, size, suffix_dim_size});
161 Eigen::DSizes<Eigen::DenseIndex, 3> slice_sizes{1, size, suffix_dim_size};
    [all...]

Completed in 169 milliseconds