[TE/JAX] Remove tuple wrapper of singleton in HLO lowering return (#1000)

* removed singleton wrapper
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

transformer_engine/jax/cpp_extensions/activation.py

@@ -98,7 +98,7 @@ class ActLuPrimitive(BasePrimitive):
 
         out = custom_caller(ActLuPrimitive.name, args, opaque, False)
 
-        return [out]
+        return out
 
     @staticmethod
     def impl(x, act_enum):
@@ -226,7 +226,7 @@ class DActLuPrimitive(BasePrimitive):
 
         out = custom_caller(DActLuPrimitive.name, args, opaque, False)
 
-        return [out]
+        return out
 
     @staticmethod
     def impl(dz, x, act_enum):

transformer_engine/jax/cpp_extensions/softmax.py

@@ -135,7 +135,7 @@ class SoftmaxPrimitive(BasePrimitive):
 
         out = custom_caller(name, args, opaque, False)
 
-        return [out]
+        return out
 
     @staticmethod
     def forward_impl(primitive, logits, scale_factor):
@@ -247,7 +247,7 @@ class SoftmaxPrimitive(BasePrimitive):
 
         out = custom_caller(name, args, opaque, False)
 
-        return [out]
+        return out
 
     @staticmethod
     def backward_impl(primitive, dz, softmax_out, scale_factor):
@@ -577,7 +577,7 @@ class ScaledMaskedSoftmaxFwdPrimitive(SoftmaxPrimitive):
 
         out = custom_caller(ScaledMaskedSoftmaxFwdPrimitive.name, args, opaque, False)
 
-        return [out]
+        return out
 
     @staticmethod
     def impl(logits, mask, scale_factor):

transformer_engine/jax/cpp_extensions/transpose.py

@@ -103,7 +103,7 @@ class TransposePrimitive(BasePrimitive):
 
         out = custom_caller(TransposePrimitive.name, args, opaque, False)
 
-        return [out]
+        return out
 
     @staticmethod
     def impl(x, static_axis_boundary, transpose_axis_boundary):