Optimizing amatrix.py

Signed-off-by: shengwang <waswang@gmail.com>