import numpy as np
import pytest

import megengine.functional as F
from megengine import tensor
from megengine.test import assertTensorClose


def test_onehot_low_dimension():
    inp = tensor(np.arange(1, 4, dtype=np.int32))
    out = 
F.one_hot(inp)