[mlir][tosa] Change axis to I32Attr

This patch changes axis attribute type from I64Attr to I32Attr
to match the TOSA spec.

Signed-off-by: Tai Ly <tai.ly@arm.com>
Change-Id: I88416f0b102f8ac2191da63a3fbce25759f19613

Reviewed By: rsuderman

Differential Revision: https://reviews.llvm.org/D157424