Merge pull request #50 from microsoft/dev/masonr/fix_last_major_rocm_alloc

Fix byte size calculation to support layouts other than FIRST_MAJOR

hatlib/arg_info.py

@@ -54,7 +54,8 @@ class ArgInfo:
         def product(l):
             return reduce(lambda x1, x2: x1*x2, l)
 
-        self.total_byte_size = self.element_num_bytes * self.numpy_shape[0] * product(self.element_strides)
+        major_dim = self.element_strides.index(max(self.element_strides))
+        self.total_byte_size = self.element_num_bytes * self.numpy_shape[major_dim] * product(self.element_strides)
 
 
 # TODO: Update this to take a HATFunction instead, instead of arg_infos and function_name