parameters: (self, x, dims)
return_statement: return lbann.Reshape(y, dims=str_list(dims))
docstring: Apply fftshift.
  
          Args:
              x (lbann.Layer): Input tensor
              dims (tuple of int): Dimensions of x (dim 0 corresponds to
                  channel)
  
          Returns:
              Layer: Output tensor
docstring_summary: Apply fftshift.
docstring_tokens: Apply fftshift .
function: def forward(self, x, dims):
          """Apply fftshift.
  
          Args:
              x (lbann.Layer): Input tensor
              dims (tuple of int): Dimensions of x (dim 0 corresponds to
                  channel)
  
          Returns:
              Layer: Output tensor
  
          """
  
          # Get gather indices by applying fftshift to tensor filled with indices
          # Note: Independent fftshift for each channel (dim 0)
          spatial_size = np.prod(dims[1:])
          spatial_inds = np.arange(spatial_size).reshape(dims[1:])
          spatial_inds = np.fft.fftshift(spatial_inds)
          channel_offsets = np.arange(0, dims[0]*spatial_size, spatial_size)
          channel_offsets = channel_offsets.reshape([-1] + [1]*spatial_inds.ndim)
          inds = np.expand_dims(spatial_inds, 0) + channel_offsets
  
          # Construct LBANN layer graph
          size = np.prod(dims)
          x = lbann.Reshape(x, dims=str_list([size]))
          inds = lbann.WeightsLayer(
              weights=lbann.Weights(
                  lbann.ValueInitializer(values=str_list(inds.flatten())),
                  optimizer=lbann.NoOptimizer(),
              ),
              dims=str_list([size]),
          )
          y = lbann.Gather(x, inds)
          return lbann.Reshape(y, dims=str_list(dims))
function_tokens: def forward ( self , x , dims ) : # Get gather indices by applying fftshift to tensor filled with indices # Note: Independent fftshift for each channel (dim 0) spatial_size = np . prod ( dims [ 1 : ] ) spatial_inds = np . arange ( spatial_size ) . reshape ( dims [ 1 : ] ) spatial_inds = np . fft . fftshift ( spatial_inds ) channel_offsets = np . arange ( 0 , dims [ 0 ] * spatial_size , spatial_size ) channel_offsets = channel_offsets . reshape ( [ - 1 ] + [ 1 ] * spatial_inds . ndim ) inds = np . expand_dims ( spatial_inds , 0 ) + channel_offsets # Construct LBANN layer graph size = np . prod ( dims ) x = lbann . Reshape ( x , dims = str_list ( [ size ] ) ) inds = lbann . WeightsLayer ( weights = lbann . Weights ( lbann . ValueInitializer ( values = str_list ( inds . flatten ( ) ) ) , optimizer = lbann . NoOptimizer ( ) , ) , dims = str_list ( [ size ] ) , ) y = lbann . Gather ( x , inds ) return lbann . Reshape ( y , dims = str_list ( dims ) )