How to use Numba extension in setup.py?

I try to use the Numba for some fast calculation. I got the following issue with create a package use Numba extension.

I did similar things as suggested on the official website. I have the following folder structure:

-test_numba
  -test_numba
    -__init__.py
    -source_module.py
-setup.py

Then I put the following code in source_module.py:

from numba.pycc import CC

cc = CC('my_module')
cc.verbose = True

@cc.export('multf', 'f8(f8, f8)')
@cc.export('multi', 'i4(i4, i4)')
def mult(a, b):
    return a * b

@cc.export('square', 'f8(f8)')
def square(a):
    return a ** 2

if __name__ == "__main__":
    cc.compile()

And complied it. and for init.py, I simple import the functions:

from .my_module import *

and for setup.py:

from setuptools import setup
from test_numba.source_module import cc
from setuptools import Extension

if __name__ == "__main__":
    setup(
        name="test_numba",
        version="0.0.1",
        packages=["test_numba"],
        ext_modules=[cc.distutils_extension()]
    )

I used pip install to installed the package. but the extension is not installed in the correct folder location.

    /Users/xxx/miniconda3/envs/py/lib/python3.7/site-packages/my_module.cpython-37m-darwin.so
    /Users/xxx/miniconda3/envs/py/lib/python3.7/site-packages/test_numba-0.0.1.dist-info/*
    /Users/xxx/miniconda3/envs/py/lib/python3.7/site-packages/test_numba/*

and when I try to import this package from jupyter. I will get the following error: No module named 'test_numba.my_module'

How to correct this?

from How to use Numba extension in setup.py?