Searched refs:output_axis (Results 1 – 2 of 2) sorted by relevance
2544 int output_axis, std::vector<uint64_t> *input_indices, in ComputePermutation() argument2548 assert(output_axis < num_dimensions); in ComputePermutation()2549 const int input_axis = perm[output_axis]; in ComputePermutation()2550 for (int i = 0; i < output_shape[output_axis]; ++i) { in ComputePermutation()2555 const bool is_last_axis = output_axis == num_dimensions - 1; in ComputePermutation()2560 output_axis + 1, input_indices, new_values); in ComputePermutation()
2362 const int output_axis = is_conv2d_backprop_input ? 1 : 0; in ConvertConv2DHelper() local2363 const int noutput = weights.shape_.d[output_axis] * num_groups; in ConvertConv2DHelper()3353 const int output_axis = is_conv3d_backprop_input ? 1 : 0; in ConvertConv3DHelper() local3354 const int noutput = weights.shape_.d[output_axis] * num_groups; in ConvertConv3DHelper()