tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc
// Returns the transposed filter for use in BackpropInput of group convolution.
Returns the transposed filter for use in BackpropInput of group convolution.
Returns the transposed filter for use in BackpropInput of group convolution .
xla::XlaOp TransposeFilterForGroupConvolutionBackpropInput(
    xla::XlaOp filter, const xla::Shape& filter_shape, int64_t num_groups,
    int num_spatial_dims) {
  int num_dims = filter_shape.dimensions_size();
  CHECK_GE(num_dims, 2);  
  xla::Shape new_shape = filter_shape;
  new_shape.set_dimensions(num_dims - 1, num_groups);
  new_shape.add_dimensions(filter_shape.dimensions(num_dims - 1) / num_groups);
  xla::XlaOp result = xla::Reshape(filter, new_shape.dimensions());
  std::vector<int64_t> transpose_dims(num_dims + 1);
  std::iota(transpose_dims.begin(), transpose_dims.end(), 0);
  std::swap(transpose_dims[num_spatial_dims],
            transpose_dims[num_spatial_dims + 1]);
  result = xla::Transpose(result, transpose_dims);
  result = xla::Collapse(result, {num_spatial_dims, num_spatial_dims + 1});
  return result;
}
xla :: XlaOp TransposeFilterForGroupConvolutionBackpropInput ( xla :: XlaOp filter , const xla :: Shape & filter_shape , int64_t num_groups , int num_spatial_dims ) { int num_dims = filter_shape . dimensions_size ( ) ; CHECK_GE ( num_dims , 2 ) ; xla :: Shape new_shape = filter_shape ; new_shape . set_dimensions ( num_dims - 1 , num_groups ) ; new_shape . add_dimensions ( filter_shape . dimensions ( num_dims - 1 ) / num_groups ) ; xla :: XlaOp result = xla :: Reshape ( filter , new_shape . dimensions ( ) ) ; std :: vector < int64_t > transpose_dims ( num_dims + 1 ) ; std :: iota ( transpose_dims . begin ( ) , transpose_dims . end ( ) , 0 ) ; std :: swap ( transpose_dims [ num_spatial_dims ] , transpose_dims [ num_spatial_dims + 1 ] ) ; result = xla :: Transpose ( result , transpose_dims ) ; result = xla :: Collapse ( result , { num_spatial_dims , num_spatial_dims + 1 } ) ; return result ; }
Returns the transposed filter for use in BackpropInput of group convolution.
Returns the transposed filter for use in BackpropInput of group convolution .
// 1. Reshape from [H, W, ..., filter_in_depth, out_depth] to [H, W, ..., // filter_in_depth, G, out_depth / G] // Crash OK // 2. Transpose to [H, W, ..., G, filter_in_depth, out_depth / G] // 3. Reshape to [H, W, ..., in_depth, out_depth / G]