test/python/dialects/shape.py

# RUN: %PYTHON %s | FileCheck %s

from mlir.ir import *
import numpy as np
import mlir.dialects.builtin as builtin
import mlir.dialects.shape as shape


def run(f):
  print("\nTEST:", f.__name__)
  f()
  return f


# CHECK-LABEL: TEST: testConstShape
@run
def testConstShape():
  with Context() as ctx, Location.unknown():
    module = Module.create()
    f32 = F32Type.get()
    with InsertionPoint(module.body):
      @builtin.FuncOp.from_py_func(
          RankedTensorType.get((12, -1), f32))
      def const_shape_tensor(arg):
        return shape.ConstShapeOp(
          DenseElementsAttr.get(np.array([10, 20]), type=IndexType.get()))

    # CHECK-LABEL: func @const_shape_tensor(%arg0: tensor<12x?xf32>)
    # CHECK: shape.const_shape [10, 20] : tensor<2xindex>
    print(module)