Function test_split_infer_struct_info

tests/python/relax/test_op_manipulate.py:2190–2296

()

def test_split_infer_struct_info():
    bb = relax.BlockBuilder()
    n = tirx.Var("n", "int64")
    x = relax.Var("x", R.Tensor((16, 4)))
    y = relax.Var("y", R.Tensor((16, 4), "float32"))
    z = relax.Var("z", R.Tensor((n, 16)))
    w = relax.Var("w", R.Tensor((n + 5, 16)))

    # All relax shape variables are non-negative.  When a scope
    # begins, any TIR variables that are used as shape variables are
    # declared to be non-negative `tvm.arith.Analyzer`.  Because
    # `relax.op.split` clamps the indices to be within the bounds of
    # the axis being split, simplifying with non-negative shape
    # variables can result in much simpler shapes.
    #
    # For example, an axis of size `n`, split on the range from 2 to 5
    # has size `T.max(T.min(5, n + 5) - T.min(2, n + 5), 0)`.  If it
    # is known that `n >= 0`, then this simplifies down to `3`.
    bb.begin_scope([x, y, z, w])

    _check_inference(
        bb,
        relax.op.split(x, 1),
        R.Tensor([16, 4]),
    )
    _check_inference(
        bb,
        relax.op.split(x, 2),
        R.Tuple(
            R.Tensor([8, 4]),
            R.Tensor([8, 4]),
        ),
    )
    # Uneven splits are allowed, with the last split being smaller than the others.
    _check_inference(
        bb,
        relax.op.split(x, 3),
        R.Tuple(
            R.Tensor([6, 4]),
            R.Tensor([6, 4]),
            R.Tensor([4, 4]),
        ),
    )

    # Dtype of result is inherited from the tensor
    _check_inference(
        bb,
        relax.op.split(y, 2),
        R.Tuple(
            R.Tensor([8, 4], "float32"),
            R.Tensor([8, 4], "float32"),
        ),
    )

    # Axis can be explicitly specified.  Otherwise, defaults to axis=0.
    _check_inference(
        bb, relax.op.split(x, [2], axis=1), R.Tuple(R.Tensor([16, 2]), R.Tensor([16, 2]))
    )