Searched refs:first_dim_size (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | dynamic_stitch_op.cc | 77 OpInputList* data_inputs, int* first_dim_size, in CheckArgsAndAllocateResult() argument
98 *first_dim_size = max_index + 1; in CheckArgsAndAllocateResult()
128 result_shape.AddDim(*first_dim_size); in CheckArgsAndAllocateResult()
140 const int32 slice_size, const int32 first_dim_size,
147 const int32 first_dim_size, \
166 int first_dim_size; in Compute() local
170 &first_dim_size, &data_elements_size, in Compute()
180 if (first_dim_size > 0) { in Compute()
188 GpuDeviceArrayOnHost<int32> indices_flat(c, first_dim_size); in Compute()
193 for (int i = 0; i < first_dim_size; ++i) { in Compute()
[all …]
|
| D | scatter_functor_gpu.cu.h | 88 Index first_dim_size, Index updates_size,
96 if (!(param_first_index >= 0 && param_first_index < first_dim_size)) {
109 Index first_dim_size,
118 if (!(param_first_index >= 0 && param_first_index < first_dim_size)) {
140 const Index first_dim_size = params.dimension(0);
147 indices.data(), first_dim_size, updates_size, indices_size));
161 const Index first_dim_size = params.dimension(0);
168 params.data(), update.data(), indices.data(), first_dim_size,
|
| D | dynamic_stitch_op_gpu.cu.cc | 53 const int32 slice_size, const int32 first_dim_size, in DynamicStitchGPUImpl() argument
57 const int32 output_size = first_dim_size * slice_size; in DynamicStitchGPUImpl()
69 const int32 first_dim_size, \
|
| D | training_ops.cc | 176 const Tindex first_dim_size = static_cast<Tindex>(var.dimension(0)); in operator ()() local
187 if (!FastBoundsCheck(index, first_dim_size)) { in operator ()()
215 if (!FastBoundsCheck(index, first_dim_size)) { in operator ()()
285 const Tindex first_dim_size = static_cast<Tindex>(var.dimension(0)); in operator ()() local
292 if (!FastBoundsCheck(index, first_dim_size)) { in operator ()()
320 if (!FastBoundsCheck(index, first_dim_size)) { in operator ()()
612 const Tindex first_dim_size = in operator ()() local
617 if (!FastBoundsCheck(index, first_dim_size)) { in operator ()()
647 const Tindex first_dim_size = accum_flat.size(); in operator ()() local
651 if (!FastBoundsCheck(index, first_dim_size)) { in operator ()()
[all …]
|
| D | training_ops_gpu.cu.cc | 512 const Tindex first_dim_size = var.dimension(0); in operator ()() local
522 lr.data(), epsilon.data(), grad.data(), indices.data(), first_dim_size, in operator ()()
573 const Tindex first_dim_size = var.dimension(0); in operator ()() local
584 first_dim_size, grad_size, indices_size); in operator ()()
783 const Tindex first_dim_size = var.dimension(0); in operator ()() local
801 /*indices=*/indices.data(), /*param_rows=*/first_dim_size, in operator ()()
859 const Tindex first_dim_size = var.dimension(0); in operator ()() local
868 first_dim_size, grad_size, indices_size)); in operator ()()
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf.cc | 4995 int64_t first_dim_size = input_type.getDimSize(0); in matchAndRewrite() local
4996 if (ShapedType::isDynamic(first_dim_size)) return failure(); in matchAndRewrite()
5000 if (first_dim_size <= 1) { in matchAndRewrite()
5069 RankedTensorType::get({first_dim_size}, rewriter.getIntegerType(32)); in matchAndRewrite()
5074 Value swaps = CreateRngUniform32(op.getLoc(), first_dim_size, 0, in matchAndRewrite()
5075 first_dim_size, &rewriter); in matchAndRewrite()
5120 CreateWhile32(op.getLoc(), first_dim_size, swap_body_fn, {swaps, indices}, in matchAndRewrite()
|