1// RUN: mlir-opt %s | mlir-opt | FileCheck %s
2
3// TODO: check lowering to an actual implementation
4
5#SparseVector = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
6#SparseMatrix = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
7
8// CHECK-LABEL: func.func @sparse_expand(
9// CHECK-SAME:  %[[A:.*]]: tensor<100xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<10x10xf64, #sparse_tensor.encoding<{{{.*}}}>>
10//      CHECK:  %[[E:.*]] = tensor.expand_shape %[[A]] {{\[\[}}0, 1]] : tensor<100xf64, #sparse_tensor.encoding<{{{.*}}}>> into tensor<10x10xf64, #sparse_tensor.encoding<{{{.*}}}>>
11//      CHECK:  return %[[E]] : tensor<10x10xf64, #sparse_tensor.encoding<{{{.*}}}>>
12func.func @sparse_expand(%arg0: tensor<100xf64, #SparseVector>) -> tensor<10x10xf64, #SparseMatrix> {
13  %0 = tensor.expand_shape %arg0 [[0, 1]] :
14    tensor<100xf64, #SparseVector> into tensor<10x10xf64, #SparseMatrix>
15  return %0 : tensor<10x10xf64, #SparseMatrix>
16}
17
18// CHECK-LABEL: func.func @sparse_collapse(
19// CHECK-SAME:  %[[A:.*]]: tensor<10x10xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<100xf64, #sparse_tensor.encoding<{{{.*}}}>>
20//      CHECK:  %[[C:.*]] = tensor.collapse_shape %[[A]] {{\[\[}}0, 1]] : tensor<10x10xf64, #sparse_tensor.encoding<{{{.*}}}>> into tensor<100xf64, #sparse_tensor.encoding<{{{.*}}}>>
21//      CHECK:  return %[[C]] : tensor<100xf64, #sparse_tensor.encoding<{{{.*}}}>>
22func.func @sparse_collapse(%arg0: tensor<10x10xf64, #SparseMatrix>) -> tensor<100xf64, #SparseVector> {
23  %0 = tensor.collapse_shape %arg0 [[0, 1]] :
24    tensor<10x10xf64, #SparseMatrix> into tensor<100xf64, #SparseVector>
25  return %0 : tensor<100xf64, #SparseVector>
26}
27