1# RUN: %PYTHON %s | FileCheck %s
2
3from mlir.ir import *
4import numpy as np
5import mlir.dialects.func as func
6import mlir.dialects.shape as shape
7
8
9def run(f):
10  print("\nTEST:", f.__name__)
11  f()
12  return f
13
14
15# CHECK-LABEL: TEST: testConstShape
16@run
17def testConstShape():
18  with Context() as ctx, Location.unknown():
19    module = Module.create()
20    f32 = F32Type.get()
21    with InsertionPoint(module.body):
22      @func.FuncOp.from_py_func(
23          RankedTensorType.get((12, -1), f32))
24      def const_shape_tensor(arg):
25        return shape.ConstShapeOp(
26          DenseElementsAttr.get(np.array([10, 20], dtype=np.int64), type=IndexType.get()))
27
28    # CHECK-LABEL: func @const_shape_tensor(%arg0: tensor<12x?xf32>)
29    # CHECK: shape.const_shape [10, 20] : tensor<2xindex>
30    print(module)
31