feat: check types in splines to allow amp (#52)

* Cast outputs to allow amp

* Only cast if necessary

normflows/utils/splines.py

@@ -71,8 +71,8 @@ def unconstrained_rational_quadratic_spline(
         top = tail_bound
 
     (
-        outputs[inside_interval_mask],
-        logabsdet[inside_interval_mask],
+        outputs_masked,
+        logabsdet_masked
     ) = rational_quadratic_spline(
         inputs=inputs[inside_interval_mask],
         unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
@@ -87,6 +87,12 @@ def unconstrained_rational_quadratic_spline(
         min_bin_height=min_bin_height,
         min_derivative=min_derivative,
     )
+    if outputs.dtype == outputs_masked.dtype and logabsdet.dtype == logabsdet_masked.dtype:
+        outputs[inside_interval_mask] = outputs_masked
+        logabsdet[inside_interval_mask] = logabsdet_masked
+    else:
+        outputs[inside_interval_mask] = outputs_masked.to(outputs.dtype)
+        logabsdet[inside_interval_mask] = logabsdet_masked.to(logabsdet.dtype)
 
     return outputs, logabsdet