tensorflow/compiler/mlir/tensorflow/transforms/tpu_space_to_depth_pass.cc
// Checks if input shape of convolution is good for space to depth transform.
Checks if input shape of convolution is good for space to depth transform.
Checks if input shape of convolution is good for space to depth transform .
bool Conv2DInputShapeCanTransform(Value input) {
  auto ranked_type = input.getType().dyn_cast<RankedTensorType>();
  if (!ranked_type) return false;
  auto input_shape = ranked_type.getShape();
  int32_t batch_size = input_shape[0];
  int32_t channel = input_shape[3];
  if (batch_size > 8 || channel > 8) {
    return false;
  }
  return true;
}
bool Conv2DInputShapeCanTransform ( Value input ) { auto ranked_type = input . getType ( ) . dyn_cast < RankedTensorType > ( ) ; if ( ! ranked_type ) return false ; auto input_shape = ranked_type . getShape ( ) ; int32_t batch_size = input_shape [ 0 ] ; int32_t channel = input_shape [ 3 ] ; if ( batch_size > 8 || channel > 8 ) { return false ; } return true ; }
Checks if input shape of convolution is good for space to depth transform.
Checks if input shape of convolution is good for space to depth transform .