/external/tensorflow/tensorflow/compiler/xla/tests/ |
D | convolution_test.cc | 254 ConvolutionDimensionNumbers dnums; in XLA_TEST_F() local 255 dnums.set_input_batch_dimension(0); in XLA_TEST_F() 256 dnums.set_output_batch_dimension(0); in XLA_TEST_F() 257 dnums.add_input_spatial_dimensions(1); in XLA_TEST_F() 258 dnums.add_output_spatial_dimensions(1); in XLA_TEST_F() 259 dnums.add_input_spatial_dimensions(2); in XLA_TEST_F() 260 dnums.add_output_spatial_dimensions(2); in XLA_TEST_F() 261 dnums.add_input_spatial_dimensions(3); in XLA_TEST_F() 262 dnums.add_output_spatial_dimensions(3); in XLA_TEST_F() 263 dnums.set_input_feature_dimension(4); in XLA_TEST_F() [all …]
|
D | convolution_variants_test.cc | 980 ConvolutionDimensionNumbers dnums; in XLA_TEST_F() local 982 dnums.set_input_batch_dimension(0); in XLA_TEST_F() 983 dnums.set_output_batch_dimension(0); in XLA_TEST_F() 984 dnums.add_input_spatial_dimensions(1); in XLA_TEST_F() 985 dnums.add_output_spatial_dimensions(1); in XLA_TEST_F() 986 dnums.add_input_spatial_dimensions(2); in XLA_TEST_F() 987 dnums.add_output_spatial_dimensions(2); in XLA_TEST_F() 988 dnums.set_input_feature_dimension(3); in XLA_TEST_F() 989 dnums.set_output_feature_dimension(3); in XLA_TEST_F() 992 dnums.add_kernel_spatial_dimensions(0); in XLA_TEST_F() [all …]
|
D | convolution_test_1d.cc | 93 ConvolutionDimensionNumbers dnums; in TestImpl() local 94 dnums.set_input_batch_dimension(0); in TestImpl() 95 dnums.set_output_batch_dimension(0); in TestImpl() 96 dnums.add_input_spatial_dimensions(1); in TestImpl() 97 dnums.add_output_spatial_dimensions(1); in TestImpl() 98 dnums.set_input_feature_dimension(2); in TestImpl() 99 dnums.set_output_feature_dimension(2); in TestImpl() 100 dnums.add_kernel_spatial_dimensions(0); in TestImpl() 101 dnums.set_kernel_input_feature_dimension(1); in TestImpl() 102 dnums.set_kernel_output_feature_dimension(2); in TestImpl() [all …]
|
/external/tensorflow/tensorflow/compiler/mlir/xla/ |
D | attribute_importer.cc | 49 const xla::GatherDimensionNumbers& dnums, mlir::Builder* builder) { in ConvertGatherDimensionNumbers() argument 50 std::vector<int64_t> offset_dims(dnums.offset_dims().begin(), in ConvertGatherDimensionNumbers() 51 dnums.offset_dims().end()); in ConvertGatherDimensionNumbers() 53 dnums.collapsed_slice_dims().begin(), dnums.collapsed_slice_dims().end()); in ConvertGatherDimensionNumbers() 54 std::vector<int64_t> start_index_map(dnums.start_index_map().begin(), in ConvertGatherDimensionNumbers() 55 dnums.start_index_map().end()); in ConvertGatherDimensionNumbers() 59 builder->getI64IntegerAttr(dnums.index_vector_dim()), in ConvertGatherDimensionNumbers() 64 const xla::ScatterDimensionNumbers& dnums, mlir::Builder* builder) { in ConvertScatterDimensionNumbers() argument 65 std::vector<int64_t> update_window_dims(dnums.update_window_dims().begin(), in ConvertScatterDimensionNumbers() 66 dnums.update_window_dims().end()); in ConvertScatterDimensionNumbers() [all …]
|
D | attribute_importer.h | 34 const xla::GatherDimensionNumbers& dnums, mlir::Builder* builder); 38 const xla::ScatterDimensionNumbers& dnums, mlir::Builder* builder); 42 const DotDimensionNumbers& dnums, mlir::Builder* builder); 46 const xla::ConvolutionDimensionNumbers& dnums, mlir::Builder* builder);
|
/external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
D | conv_canonicalization_test.cc | 70 ConvolutionDimensionNumbers dnums; in TEST_F() local 71 dnums.set_input_batch_dimension(1); in TEST_F() 72 dnums.set_output_batch_dimension(1); in TEST_F() 73 dnums.add_input_spatial_dimensions(2); in TEST_F() 74 dnums.add_output_spatial_dimensions(2); in TEST_F() 75 dnums.add_input_spatial_dimensions(3); in TEST_F() 76 dnums.add_output_spatial_dimensions(3); in TEST_F() 77 dnums.set_input_feature_dimension(0); in TEST_F() 78 dnums.set_output_feature_dimension(0); in TEST_F() 79 dnums.add_kernel_spatial_dimensions(2); in TEST_F() [all …]
|
D | ir_emission_utils.cc | 83 const ConvolutionDimensionNumbers& dnums = in PotentiallyImplementedAsEigenConvolution() local 87 const int64 num_spatial_dims = dnums.output_spatial_dimensions_size(); in PotentiallyImplementedAsEigenConvolution() 93 if (dnums.input_spatial_dimensions(i) != i + 1) { in PotentiallyImplementedAsEigenConvolution() 96 if (dnums.kernel_spatial_dimensions(i) != i) { in PotentiallyImplementedAsEigenConvolution() 99 if (dnums.output_spatial_dimensions(i) != i + 1) { in PotentiallyImplementedAsEigenConvolution() 104 return dnums.input_batch_dimension() == 0 && in PotentiallyImplementedAsEigenConvolution() 105 dnums.input_feature_dimension() == input_shape.dimensions_size() - 1 && in PotentiallyImplementedAsEigenConvolution() 106 dnums.output_batch_dimension() == 0 && in PotentiallyImplementedAsEigenConvolution() 107 dnums.output_feature_dimension() == in PotentiallyImplementedAsEigenConvolution() 109 dnums.kernel_input_feature_dimension() == in PotentiallyImplementedAsEigenConvolution() [all …]
|
D | conv_canonicalization.cc | 39 const ConvolutionDimensionNumbers& dnums = in Run() local 41 auto input_batch_dim = dnums.input_batch_dimension(); in Run() 42 auto input_feature_dim = dnums.input_feature_dimension(); in Run() 43 auto kernel_input_feature_dim = dnums.kernel_input_feature_dimension(); in Run() 44 auto kernel_output_feature_dim = dnums.kernel_output_feature_dimension(); in Run() 46 const int64 num_spatial_dims = dnums.output_spatial_dimensions_size(); in Run() 65 new_input_dim_order[i + 1] = dnums.input_spatial_dimensions(i); in Run() 67 input->shape().dimensions(dnums.input_spatial_dimensions(i)); in Run() 84 new_kernel_dim_order[i] = dnums.kernel_spatial_dimensions(i); in Run() 86 kernel->shape().dimensions(dnums.kernel_spatial_dimensions(i)); in Run() [all …]
|
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
D | einsum.cc | 139 EinsumDimensionNumbers dnums; in GetEinsumDimensionNumbers() local 144 dnums.lhs.emplace_back(i); in GetEinsumDimensionNumbers() 146 dnums.lhs_out.emplace_back(i, out_index->second); in GetEinsumDimensionNumbers() 148 dnums.lhs_rhs.emplace_back(i, rhs_index->second); in GetEinsumDimensionNumbers() 150 dnums.lhs_rhs_out.emplace_back(i, rhs_index->second, out_index->second); in GetEinsumDimensionNumbers() 158 dnums.rhs.emplace_back(i); in GetEinsumDimensionNumbers() 160 dnums.rhs_out.emplace_back(i, out_index->second); in GetEinsumDimensionNumbers() 172 return dnums; in GetEinsumDimensionNumbers() 187 const Location& loc, EinsumDimensionNumbers& dnums, Value* lhs, Value* rhs, in transposeForBatchMatmul() argument 192 lhs_transpose.reserve(dnums.lhs_rhs_out.size() + dnums.lhs_out.size() + in transposeForBatchMatmul() [all …]
|
/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
D | convolution_handler.cc | 55 const auto& dnums = original_hlo->convolution_dimension_numbers(); in PartitionConvolutionWithBatchGroupCount() local 58 lhs.base_shape().dimensions(dnums.input_batch_dimension()); in PartitionConvolutionWithBatchGroupCount() 60 rhs.base_shape().dimensions(dnums.kernel_output_feature_dimension()); in PartitionConvolutionWithBatchGroupCount() 68 rhs_to_lhs_indices[dnums.kernel_output_feature_dimension()] = in PartitionConvolutionWithBatchGroupCount() 69 dnums.input_batch_dimension(); in PartitionConvolutionWithBatchGroupCount() 70 rhs_to_lhs_indices[dnums.kernel_input_feature_dimension()] = in PartitionConvolutionWithBatchGroupCount() 71 dnums.input_feature_dimension(); in PartitionConvolutionWithBatchGroupCount() 72 for (int64 i = 0; i < dnums.input_spatial_dimensions_size(); ++i) { in PartitionConvolutionWithBatchGroupCount() 73 rhs_to_lhs_indices[dnums.kernel_spatial_dimensions(i)] = in PartitionConvolutionWithBatchGroupCount() 74 dnums.input_spatial_dimensions(i); in PartitionConvolutionWithBatchGroupCount() [all …]
|
D | gather_scatter_handler.cc | 139 GatherDimensionNumbers dnums = gather->gather_dimension_numbers(); in PartitionIndexOnlyPartition() local 142 (dnums.index_vector_dim() == indices.base_shape().rank() || in PartitionIndexOnlyPartition() 143 indices.sharding().tile_assignment().dim(dnums.index_vector_dim()) == in PartitionIndexOnlyPartition() 149 indices.hlo()->shape(), dnums, in PartitionIndexOnlyPartition() 157 int64 indices_batch_dim = i < dnums.index_vector_dim() ? i : i + 1; in PartitionIndexOnlyPartition() 185 GatherDimensionNumbers dnums = gather->gather_dimension_numbers(); in ParititonPassthroughOperand() local 200 pshape, operand.hlo(), indices.hlo(), dnums, pslice_sizes, in ParititonPassthroughOperand() 220 GatherDimensionNumbers dnums = gather->gather_dimension_numbers(); in ParititonTrivialIndexedOperandDimension() local 221 std::vector<int64> start_index_map(dnums.start_index_map().begin(), in ParititonTrivialIndexedOperandDimension() 222 dnums.start_index_map().end()); in ParititonTrivialIndexedOperandDimension() [all …]
|
/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
D | stream_executor_util.cc | 49 StreamExecutorConvLayoutsToXlaLayouts(const ConvolutionDimensionNumbers& dnums, in StreamExecutorConvLayoutsToXlaLayouts() argument 55 input_layout.push_back(dnums.input_batch_dimension()); in StreamExecutorConvLayoutsToXlaLayouts() 56 input_layout.push_back(dnums.input_feature_dimension()); in StreamExecutorConvLayoutsToXlaLayouts() 58 dnums.input_spatial_dimensions().begin(), in StreamExecutorConvLayoutsToXlaLayouts() 59 dnums.input_spatial_dimensions().end()); in StreamExecutorConvLayoutsToXlaLayouts() 62 input_layout.push_back(dnums.input_batch_dimension()); in StreamExecutorConvLayoutsToXlaLayouts() 64 dnums.input_spatial_dimensions().begin(), in StreamExecutorConvLayoutsToXlaLayouts() 65 dnums.input_spatial_dimensions().end()); in StreamExecutorConvLayoutsToXlaLayouts() 66 input_layout.push_back(dnums.input_feature_dimension()); in StreamExecutorConvLayoutsToXlaLayouts() 71 ConvolutionDimensionNumbersToString(dnums)); in StreamExecutorConvLayoutsToXlaLayouts() [all …]
|
D | cudnn_pad_for_convolutions.cc | 182 const auto& dnums = conv->convolution_dimension_numbers(); in TryResolvePaddedShapesForTensorCore() local 227 new_input_shape->dimensions(dnums.input_feature_dimension()); in TryResolvePaddedShapesForTensorCore() 229 new_output_shape->dimensions(dnums.output_feature_dimension()); in TryResolvePaddedShapesForTensorCore() 231 new_input_shape->set_dimensions(dnums.input_feature_dimension(), 4); in TryResolvePaddedShapesForTensorCore() 232 new_filter_shape->set_dimensions(dnums.kernel_input_feature_dimension(), 4); in TryResolvePaddedShapesForTensorCore() 237 pad_dim(new_input_shape, dnums.input_feature_dimension()); in TryResolvePaddedShapesForTensorCore() 238 pad_dim(new_filter_shape, dnums.kernel_input_feature_dimension()); in TryResolvePaddedShapesForTensorCore() 239 pad_dim(new_filter_shape, dnums.kernel_output_feature_dimension()); in TryResolvePaddedShapesForTensorCore() 240 pad_dim(new_output_shape, dnums.output_feature_dimension()); in TryResolvePaddedShapesForTensorCore() 299 const auto& dnums = conv->convolution_dimension_numbers(); in TryResolvePaddedShapesForIntegerConvolution() local [all …]
|
D | gpu_conv_rewriter.cc | 43 const ConvolutionDimensionNumbers& dnums, in CreateGpuConv() argument 62 custom_call->set_convolution_dimension_numbers(dnums); in CreateGpuConv() 122 const ConvolutionDimensionNumbers& dnums = in CanImplementAsGpuForwardConv() local 124 if (dnums.input_spatial_dimensions_size() > 3) { in CanImplementAsGpuForwardConv() 136 if (dnums.input_spatial_dimensions_size() == 2 in CanImplementAsGpuForwardConv() 344 ConvolutionDimensionNumbers dnums = conv->convolution_dimension_numbers(); in MatchBackwardInput() local 364 absl::c_is_permutation(dnums.kernel_spatial_dimensions(), in MatchBackwardInput() 398 const auto& input_spatial_dims = dnums.input_spatial_dimensions(); in MatchBackwardInput() 399 const auto& output_spatial_dims = dnums.output_spatial_dimensions(); in MatchBackwardInput() 505 dnums.set_kernel_input_feature_dimension( in MatchBackwardInput() [all …]
|
D | gpu_conv_runner.cc | 314 const ConvolutionDimensionNumbers& dnums = desc.dnums; in GetGpuConvConfig() local 328 VLOG(3) << "Dim nums: { " << dnums.ShortDebugString() << " }"; in GetGpuConvConfig() 344 CHECK_EQ(num_dimensions, dnums.input_spatial_dimensions_size()) in GetGpuConvConfig() 346 CHECK_EQ(num_dimensions, dnums.kernel_spatial_dimensions_size()) in GetGpuConvConfig() 348 CHECK_EQ(num_dimensions, dnums.output_spatial_dimensions_size()) in GetGpuConvConfig() 371 dnums, input_shape.layout(), filter_shape.layout(), in GetGpuConvConfig() 378 input_shape.dimensions(dnums.input_feature_dimension())) in GetGpuConvConfig() 379 .set_count(input_shape.dimensions(dnums.input_batch_dimension())); in GetGpuConvConfig() 384 input_shape.dimensions(dnums.input_spatial_dimensions(dim))); in GetGpuConvConfig() 391 filter_shape.dimensions(dnums.kernel_input_feature_dimension())) in GetGpuConvConfig() [all …]
|
/external/tensorflow/tensorflow/compiler/xla/service/ |
D | dot_as_convolution_util.cc | 157 dot_as_convolution_util::DotConvolutionDimsInfo dnums; in ParseDotGeneralFromDot() local 159 dnums.batch_dims.emplace_back(); in ParseDotGeneralFromDot() 160 dnums.batch_dims.back().lhs = dot_dim_numbs.lhs_batch_dimensions(i); in ParseDotGeneralFromDot() 161 dnums.batch_dims.back().rhs = dot_dim_numbs.rhs_batch_dimensions(i); in ParseDotGeneralFromDot() 162 dnums.batch_dims.back().output = i; in ParseDotGeneralFromDot() 163 dnums.batch_dims.back().spatial_dim = -1; in ParseDotGeneralFromDot() 167 dnums.contracting_dims.emplace_back(); in ParseDotGeneralFromDot() 168 dnums.contracting_dims.back().lhs = in ParseDotGeneralFromDot() 170 dnums.contracting_dims.back().rhs = in ParseDotGeneralFromDot() 172 dnums.contracting_dims.back().output = -1; in ParseDotGeneralFromDot() [all …]
|
D | transpose_folding_test.cc | 231 auto dnums = XlaBuilder::CreateDefaultConvDimensionNumbers(); in TEST_F() local 241 transpose_y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); in TEST_F() 245 /*batch_group_count=*/1, window, dnums, in TEST_F() 250 /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums, in TEST_F() 268 EXPECT_EQ(dnums.kernel_input_feature_dimension(), in TEST_F() 271 EXPECT_EQ(dnums.kernel_output_feature_dimension(), in TEST_F() 288 auto dnums = XlaBuilder::CreateDefaultConvDimensionNumbers(); in TEST_F() local 298 transpose_y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); in TEST_F() 302 /*batch_group_count=*/1, window, dnums, in TEST_F() 307 /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums, in TEST_F() [all …]
|
D | hlo_sharding_util.cc | 505 const GatherDimensionNumbers& dnums = hlo->gather_dimension_numbers(); in GatherOutputSharding() local 508 if (absl::c_binary_search(dnums.offset_dims(), i)) { in GatherOutputSharding() 512 index_dim >= dnums.index_vector_dim() ? index_dim + 1 : index_dim; in GatherOutputSharding() 544 const GatherDimensionNumbers& dnums = hlo->gather_dimension_numbers(); in GatherIndexSharding() local 547 if (!absl::c_binary_search(dnums.offset_dims(), i)) { in GatherIndexSharding() 557 index_tile_assignment_dims.begin() + dnums.index_vector_dim(), 1); in GatherIndexSharding() 593 const GatherDimensionNumbers& dnums = hlo.gather_dimension_numbers(); in GatherEffectiveOutputSharding() local 597 if (!absl::c_binary_search(dnums.offset_dims(), i)) { in GatherEffectiveOutputSharding() 628 if (!absl::c_binary_search(dnums.offset_dims(), i)) { in GatherEffectiveOutputSharding() 645 const ScatterDimensionNumbers& dnums = hlo->scatter_dimension_numbers(); in ScatterIndexSharding() local [all …]
|
D | dot_decomposer.cc | 180 const DotDimensionNumbers& dnums = instruction->dot_dimension_numbers(); in Run() local 183 if (dnums.lhs_contracting_dimensions_size() != 1) { in Run() 189 if (dnums.lhs_batch_dimensions_size() + 2 < in Run() 191 dnums.rhs_batch_dimensions_size() + 2 < in Run() 196 if (dnums.lhs_batch_dimensions().empty() && in Run() 197 dnums.lhs_contracting_dimensions().empty()) { in Run() 201 if (dnums.lhs_batch_dimensions().empty()) { in Run() 205 dnums.lhs_batch_dimensions_size()); in Run() 207 if (!absl::c_equal(dnums.lhs_batch_dimensions(), canonical_batch_dims) || in Run() 208 !absl::c_equal(dnums.rhs_batch_dimensions(), canonical_batch_dims)) { in Run()
|
D | shape_inference_test.cc | 403 ConvolutionDimensionNumbers dnums; in TEST_F() local 407 dnums.set_input_batch_dimension(0); in TEST_F() 408 dnums.set_output_batch_dimension(0); in TEST_F() 409 dnums.set_input_feature_dimension(1); in TEST_F() 410 dnums.set_output_feature_dimension(1); in TEST_F() 411 dnums.add_input_spatial_dimensions(2); in TEST_F() 412 dnums.add_output_spatial_dimensions(2); in TEST_F() 413 dnums.add_input_spatial_dimensions(3); in TEST_F() 414 dnums.add_output_spatial_dimensions(3); in TEST_F() 418 dnums.set_kernel_input_feature_dimension(2); in TEST_F() [all …]
|
D | sharding_propagation.cc | 132 const auto& dnums = instruction->convolution_dimension_numbers(); in IsConvolutionKernelSmall() local 133 for (int64 i = 0; i < dnums.input_spatial_dimensions().size(); ++i) { in IsConvolutionKernelSmall() 135 rhs->shape().dimensions(dnums.kernel_spatial_dimensions(i)); in IsConvolutionKernelSmall() 137 instruction->shape().dimensions(dnums.output_spatial_dimensions(i)); in IsConvolutionKernelSmall() 338 const dot_as_convolution_util::DotConvolutionDimsInfo& dnums, in InferDotShardingFromOperands() argument 347 contracting_dims.reserve(dnums.contracting_dims.size()); in InferDotShardingFromOperands() 348 for (const auto& dim : dnums.contracting_dims) { in InferDotShardingFromOperands() 354 ? dnums.rhs_non_contracting_dims in InferDotShardingFromOperands() 355 : dnums.lhs_non_contracting_dims) { in InferDotShardingFromOperands() 366 for (const auto& dim : dnums.batch_dims) { in InferDotShardingFromOperands() [all …]
|
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
D | conv_op_helpers.cc | 327 xla::ConvolutionDimensionNumbers dnums; in MakeXlaBackpropInputConvOp() local 328 dnums.set_input_batch_dimension(batch_dim); in MakeXlaBackpropInputConvOp() 329 dnums.set_output_batch_dimension(batch_dim); in MakeXlaBackpropInputConvOp() 330 dnums.set_input_feature_dimension(feature_dim); in MakeXlaBackpropInputConvOp() 331 dnums.set_output_feature_dimension(feature_dim); in MakeXlaBackpropInputConvOp() 335 dnums.set_kernel_input_feature_dimension(attrs.num_spatial_dims + 1); in MakeXlaBackpropInputConvOp() 336 dnums.set_kernel_output_feature_dimension(attrs.num_spatial_dims); in MakeXlaBackpropInputConvOp() 356 dnums.add_input_spatial_dimensions(dim); in MakeXlaBackpropInputConvOp() 357 dnums.add_kernel_spatial_dimensions(i); in MakeXlaBackpropInputConvOp() 358 dnums.add_output_spatial_dimensions(dim); in MakeXlaBackpropInputConvOp() [all …]
|
D | reverse_sequence_op.cc | 100 xla::GatherDimensionNumbers dnums; in Compile() local 101 dnums.set_index_vector_dim(2); in Compile() 104 dnums.add_start_index_map(batch_dim_); in Compile() 105 dnums.add_start_index_map(seq_dim_); in Compile() 111 dnums.add_offset_dims(i); in Compile() 113 dnums.add_collapsed_slice_dims(i); in Compile() 122 xla::Gather(input, start_indices, dnums, slice_sizes)); in Compile()
|
/external/tensorflow/tensorflow/compiler/xla/client/ |
D | xla_builder_test.cc | 783 ConvolutionDimensionNumbers dnums; in TEST_F() local 784 dnums.set_input_batch_dimension(0); in TEST_F() 785 dnums.set_output_batch_dimension(0); in TEST_F() 786 dnums.add_input_spatial_dimensions(1); in TEST_F() 787 dnums.add_output_spatial_dimensions(1); in TEST_F() 788 dnums.add_input_spatial_dimensions(2); in TEST_F() 789 dnums.add_output_spatial_dimensions(2); in TEST_F() 790 dnums.set_input_feature_dimension(3); in TEST_F() 791 dnums.set_output_feature_dimension(3); in TEST_F() 792 dnums.add_kernel_spatial_dimensions(0); in TEST_F() [all …]
|
/external/tensorflow/tensorflow/compiler/tests/ |
D | xla_ops_test.py | 136 dnums = xla_data_pb2.ConvolutionDimensionNumbers() 138 dnums.input_batch_dimension = 0 139 dnums.input_feature_dimension = 1 140 dnums.output_batch_dimension = 0 141 dnums.output_feature_dimension = 1 142 dnums.kernel_output_feature_dimension = 0 143 dnums.kernel_input_feature_dimension = 1 144 dnums.input_spatial_dimensions.extend(range(2, 2 + num_spatial_dims)) 145 dnums.kernel_spatial_dimensions.extend(range(2, 2 + num_spatial_dims)) 146 dnums.output_spatial_dimensions.extend(range(2, 2 + num_spatial_dims)) [all …]
|