diff --git a/reference-implementation/include/emitc/tosa.h b/reference-implementation/include/emitc/tosa.h
index 2019dfba..787d62a3 100644
--- a/reference-implementation/include/emitc/tosa.h
+++ b/reference-implementation/include/emitc/tosa.h
@@ -646,7 +646,7 @@ namespace {
 // Common reduce function used by specialized TOSA reduce ops.
 template <typename Dest, typename Src, typename Computation>
 inline Dest reduce(Src operand, typename get_element_type<Src>::type initValue,
-                   int64_t dimension, Computation computation) {
+                   int32_t dimension, Computation computation) {
   static_assert(is_tensor<Src>::value, "Expected tensor argument");
   static_assert(is_tensor<Dest>::value, "Expected tensor result");
 
@@ -688,7 +688,7 @@ inline Dest reduce(Src operand, typename get_element_type<Src>::type initValue,
 
 // ArgMaxOp
 template <typename Dest, typename Src>
-inline Dest argmax(Src operand, int64_t dimension) {
+inline Dest argmax(Src operand, int32_t dimension) {
   static_assert(is_tensor<Src>::value, "Expected tensor argument");
   static_assert(is_tensor<Dest>::value, "Expected tensor result");
 
@@ -732,7 +732,7 @@ inline Dest argmax(Src operand, int64_t dimension) {
 
 // ReduceAllOp
 template <typename Dest, typename Src>
-inline Dest reduce_all(Src input, int64_t dimension) {
+inline Dest reduce_all(Src input, int32_t dimension) {
   // ReduceAllOp takes only tensors with datatype bool according to the
   // TOSA specifications.
   using ET_Src = typename get_element_type<Src>::type;
@@ -750,7 +750,7 @@ inline Dest reduce_all(Src input, int64_t dimension) {
 
 // ReduceAnyOp
 template <typename Dest, typename Src>
-inline Dest reduce_any(Src input, int64_t dimension) {
+inline Dest reduce_any(Src input, int32_t dimension) {
   // ReduceAnyOp takes only tensors with datatype bool according to the
   // TOSA specifications.
   using ET_Src = typename get_element_type<Src>::type;
@@ -768,7 +768,7 @@ inline Dest reduce_any(Src input, int64_t dimension) {
 
 // ReduceMaxOp
 template <typename Dest, typename Src>
-inline Dest reduce_max(Src input, int64_t dimension) {
+inline Dest reduce_max(Src input, int32_t dimension) {
   using ET_Src = typename get_element_type<Src>::type;
 
   auto f =
@@ -780,7 +780,7 @@ inline Dest reduce_max(Src input, int64_t dimension) {
 
 // ReduceMinOp
 template <typename Dest, typename Src>
-inline Dest reduce_min(Src input, int64_t dimension) {
+inline Dest reduce_min(Src input, int32_t dimension) {
   using ET_Src = typename get_element_type<Src>::type;
 
   auto f =
@@ -792,7 +792,7 @@ inline Dest reduce_min(Src input, int64_t dimension) {
 
 // ReduceProdOp
 template <typename Dest, typename Src>
-inline Dest reduce_prod(Src input, int64_t dimension) {
+inline Dest reduce_prod(Src input, int32_t dimension) {
   using ET_Src = typename get_element_type<Src>::type;
 
   return tosa::reduce<Dest, Src>(input, 1, dimension,
@@ -801,7 +801,7 @@ inline Dest reduce_prod(Src input, int64_t dimension) {
 
 // ReduceSumOp
 template <typename Dest, typename Src>
-inline Dest reduce_sum(Src input, int64_t dimension) {
+inline Dest reduce_sum(Src input, int32_t dimension) {
   using ET_Src = typename get_element_type<Src>::type;
 
   return tosa::reduce<Dest, Src>(input, 0, dimension, std::plus<ET_Src>{});
diff --git a/reference-implementation/unittests/tosa.cpp b/reference-implementation/unittests/tosa.cpp
index 8f0a148a..6014e688 100644
--- a/reference-implementation/unittests/tosa.cpp
+++ b/reference-implementation/unittests/tosa.cpp
@@ -1136,7 +1136,7 @@ TEST(tosa, reduce_prod) {
 TEST(tosa, reduce_sum) {
   {
     Tensor<int32_t, 2, 3> input{1, 2, 3, 4, 5, 6};
-    int64_t dimension = 0;
+    int32_t dimension = 0;
     Tensor<int32_t, 3> expected_result{5, 7, 9};
     Tensor<int32_t, 3> result =
         tosa::reduce_sum<Tensor<int32_t, 3>>(input, dimension);
@@ -1145,7 +1145,7 @@ TEST(tosa, reduce_sum) {
   }
   {
     Tensor<int32_t, 2, 3> input{1, 2, 3, 4, 5, 6};
-    int64_t dimension = 1;
+    int32_t dimension = 1;
     Tensor<int32_t, 2> expected_result{6, 15};
     Tensor<int32_t, 2> result =
         tosa::reduce_sum<Tensor<int32_t, 2>>(input, dimension);
@@ -1155,7 +1155,7 @@ TEST(tosa, reduce_sum) {
   {
     Tensor<int32_t, 4, 2, 3> input{1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6,
                                    1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6};
-    int64_t dimension = 0;
+    int32_t dimension = 0;
     Tensor<int32_t, 2, 3> expected_result{4, 8, 12, 16, 20, 24};
     Tensor<int32_t, 2, 3> result =
         tosa::reduce_sum<Tensor<int32_t, 2, 3>>(input, dimension);
@@ -1165,7 +1165,7 @@ TEST(tosa, reduce_sum) {
   {
     Tensor<int32_t, 4, 2, 3> input{1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6,
                                    1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6};
-    int64_t dimension = 1;
+    int32_t dimension = 1;
     Tensor<int32_t, 4, 3> expected_result{5, 7, 9, 5, 7, 9, 5, 7, 9, 5, 7, 9};
     Tensor<int32_t, 4, 3> result =
         tosa::reduce_sum<Tensor<int32_t, 4, 3>>(input, dimension);
@@ -1175,7 +1175,7 @@ TEST(tosa, reduce_sum) {
   {
     Tensor<int32_t, 4, 2, 3> input{1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6,
                                    1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6};
-    int64_t dimension = 2;
+    int32_t dimension = 2;
     Tensor<int32_t, 4, 2> expected_result{6, 15, 6, 15, 6, 15, 6, 15};
     Tensor<int32_t, 4, 2> result =
         tosa::reduce_sum<Tensor<int32_t, 4, 2>>(input, dimension);