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

  /external/tensorflow/tensorflow/core/kernels/
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});
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...]
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 204 milliseconds