diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst
index 7833a28c02..1e76e1be11 100644
--- a/docs/source/run_card.rst
+++ b/docs/source/run_card.rst
@@ -589,3 +589,8 @@ function. The activations are ``relu``, ``leaky_relu``, ``elu``, ``gelu``, ``sig
 
 ``cwnet_hidden_dim`` (``"auto"``), ``cwnet_layers`` (``3``), ``cwnet_activation`` (``"leaky_relu"``)
     Size and activation of the channel weight network.
+
+``nn_float_precision`` (``"float64"``)
+    Floating-point precision of the weights and the evaluation of all networks,
+    ``float32`` or ``float64``. The rest of the computation, including the network
+    outputs, softmaxes and Jacobians, stays in double precision.
diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py
index d79758a271..658cc025fd 100644
--- a/madgraph/iolibs/template_files/mg7/launch.py
+++ b/madgraph/iolibs/template_files/mg7/launch.py
@@ -3056,6 +3056,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace:
                 subnet_layers=self.madnis_settings["flow_layers"],
                 subnet_activation=self.activation(madnis_args["flow_activation"]),
                 invert_spline=madnis_args["flow_invert_spline"],
+                subnet_precision=self.nn_precision(),
             )
             if channel.adaptive_mapping is None:
                 flow.initialize_globals(self.process.contexts[0], seed)
@@ -3078,6 +3079,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace:
                     subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"],
                     subnet_layers=madnis_args["discrete_layers"],
                     subnet_activation=self.activation(madnis_args["discrete_activation"]),
+                    subnet_precision=self.nn_precision(),
                 )
                 discrete_sym.initialize_globals(self.process.contexts[0], seed)
                 cond_dim += perm_count
@@ -3092,6 +3094,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace:
                     subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"],
                     subnet_layers=madnis_args["discrete_layers"],
                     subnet_activation=self.activation(madnis_args["discrete_activation"]),
+                    subnet_precision=self.nn_precision(),
                 )
                 discrete_flavor.initialize_globals(self.process.contexts[0], seed)
 
@@ -3172,6 +3175,7 @@ def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork:
             layers=madnis_args["cwnet_layers"],
             activation=self.activation(madnis_args["cwnet_activation"]),
             prefix=f"subproc{self.subproc_id}.cwnet",
+            precision=self.nn_precision(),
         )
         cwnet.initialize_globals(
             self.process.contexts[0], self.process.run_seed
@@ -3189,6 +3193,11 @@ def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode:
         else:
             raise ValueError(f"Invalid t-channel mode '{name}'")
 
+    def nn_precision(self) -> ms.DataType:
+        # one precision for all networks: the optimizer trains all of them together
+        # the run card only allows float32 and float64
+        return getattr(ms.DataType, self.process.run_card["madnis"]["nn_float_precision"])
+
     def activation(self, name: str) -> ms.MLP.Activation:
         activations = {
             "relu": ms.MLP.relu,
diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml
index e8bc6d1fef..104bcec9cc 100644
--- a/madgraph/iolibs/template_files/mg7/run_card.toml
+++ b/madgraph/iolibs/template_files/mg7/run_card.toml
@@ -160,6 +160,9 @@ cwnet_hidden_dim = %(madnis.cwnet_hidden_dim)s
 cwnet_layers = %(madnis.cwnet_layers)s
 cwnet_activation = %(madnis.cwnet_activation)s # options: relu, leaky_relu, elu, gelu, sigmoid, softplus
 
+# precision of all networks
+nn_float_precision = %(madnis.nn_float_precision)s # options: float32, float64
+
 # training parameters
 train_batches = %(madnis.train_batches)s
 log_interval = %(madnis.log_interval)s
diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py
index 34a8890d63..597f69e5fd 100755
--- a/madgraph/various/banner.py
+++ b/madgraph/various/banner.py
@@ -6939,6 +6939,9 @@ def default_setup(self):
         self.add_toml_param('madnis', 'cwnet_activation', "leaky_relu",
             allowed=['relu', 'leaky_relu', 'elu', 'gelu', 'sigmoid', 'softplus'],
             comment="activation function of the channel weight network")
+        self.add_toml_param('madnis', 'nn_float_precision', "float64",
+            allowed=['float32', 'float64'],
+            comment="floating-point precision of the weights and the evaluation of all networks")
         self.add_toml_param('madnis', 'train_batches', 1000,
             comment="number of training batches")
         self.add_toml_param('madnis', 'log_interval', 100,
diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt
index cb663c65c3..4af576dc24 100644
--- a/madspace/CMakeLists.txt
+++ b/madspace/CMakeLists.txt
@@ -132,7 +132,7 @@ if(ENABLE_OPENBLAS)
             -DCMAKE_INSTALL_LIBDIR=lib
             -DBUILD_SHARED_LIBS=OFF
             -DBUILD_WITHOUT_LAPACK=ON
-            -DBUILD_SINGLE=OFF
+            -DBUILD_SINGLE=ON
             -DBUILD_DOUBLE=ON
             -DBUILD_COMPLEX=OFF
             -DBUILD_COMPLEX16=OFF
diff --git a/madspace/generate_headers.py b/madspace/generate_headers.py
index d77b66ba10..e379291264 100644
--- a/madspace/generate_headers.py
+++ b/madspace/generate_headers.py
@@ -25,6 +25,7 @@ def main():
     commands = {}
     for sec in data:
         commands.update({key: value for key, value in sec.items() if key != "title"})
+    commands = expand_precisions(commands)
     for i, cmd in enumerate(commands.values()):
         cmd["opcode"] = i
     sections = [
@@ -45,6 +46,51 @@ def main():
     kernel_checks(commands, out_dir)
 
 
+# C++ element type and kernel policy suffix for each float precision
+PRECISIONS = {"float32": ("float", "32"), "float64": ("double", "")}
+
+
+def expand_precisions(commands):
+    """Split each instruction with a `precision` list into one instruction per
+    precision, named `<name>_<precision>`. In these, `float` in a type means the
+    instruction's precision. The expanded entries keep the original name in
+    `base_name` and the kernel policy suffix in `types_suffix`."""
+    expanded = {}
+    for name, cmd in commands.items():
+        precisions = cmd.get("precision")
+        if precisions is None:
+            cmd["base_name"] = name
+            cmd["types_suffix"] = ""
+            cmd["dispatch"] = None
+            expanded[name] = cmd
+            continue
+        for precision in precisions:
+            ctype, suffix = PRECISIONS[precision]
+            sub_cmd = dict(cmd)
+            for key in ("inputs", "outputs"):
+                if sub_cmd[key] == "any":
+                    raise ValueError(f"{name}: precision needs explicit {key}")
+                sub_cmd[key] = [
+                    {
+                        **arg,
+                        "type": (
+                            [precision, *arg["type"][1:]]
+                            if arg.get("type", [None])[0] == "float"
+                            else arg.get("type")
+                        ),
+                    }
+                    for arg in cmd[key]
+                ]
+            sub_cmd["base_name"] = name
+            sub_cmd["types_suffix"] = suffix
+            sub_cmd["ctype"] = ctype
+            sub_cmd["doc_cmd"] = cmd
+            # the builder method is generated once, for the first precision
+            sub_cmd["dispatch"] = precisions if precision == precisions[0] else []
+            expanded[f"{name}_{precision}"] = sub_cmd
+    return expanded
+
+
 def write_autogen(f):
     f.write(
         "// This file was automatically generated based on instruction_set.yaml\n"
@@ -118,14 +164,23 @@ def function_builder_mixin(commands, out_dir):
         write_autogen(f)
         first = True
         for name, cmd in commands.items():
+            if cmd["dispatch"] == []:
+                continue
             if first:
                 first = False
             else:
                 f.write("\n")
-            f.write(doc_comment(name, cmd))
+            base_name = cmd["base_name"]
+            f.write(doc_comment(base_name, cmd.get("doc_cmd", cmd)))
             if cmd["inputs"] == "any":
                 parameters = "ValueVec args"
                 instruction_call = f'instruction("{name}", args)'
+            elif cmd["dispatch"]:
+                parameters = ", ".join(f"Value {arg['name']}" for arg in cmd["inputs"])
+                arguments = ", ".join(arg["name"] for arg in cmd["inputs"])
+                instruction_call = (
+                    f'precision_instruction("{base_name}", {{{arguments}}})'
+                )
             else:
                 parameters = ", ".join(f"Value {arg['name']}" for arg in cmd["inputs"])
                 arguments = ", ".join(arg["name"] for arg in cmd["inputs"])
@@ -152,7 +207,7 @@ def function_builder_mixin(commands, out_dir):
                         f"    return {{{return_array}}};"
                     )
 
-            f.write(f"{return_type} {name}({parameters}) {{\n{func_body}\n}}\n")
+            f.write(f"{return_type} {base_name}({parameters}) {{\n{func_body}\n}}\n")
 
 
 def instruction_set_python(commands, out_dir):
@@ -170,6 +225,9 @@ def instruction_set_python(commands, out_dir):
         )
 
         for name, cmd in commands.items():
+            if cmd["dispatch"] == []:
+                continue
+            name = cmd["base_name"]
             f.write(f'    fb.def("{name}", &FunctionBuilder::{name}')
             if cmd["inputs"] == "any":
                 f.write(', py::arg("args")')
@@ -181,12 +239,16 @@ def instruction_set_python(commands, out_dir):
         f.write("}\n}\n")
 
 
+# yaml shorthands for the default element types
+DTYPE_NAMES = {"float": "float64", "int": "int32"}
+
+
 def format_type(data):
     type = data["type"]
     if type[0] == "size":
-        return f'{{DataType::dt_int, true, {{"{type[1]}"}}, true}}'
+        return f'{{DataType::dt_int32, true, {{"{type[1]}"}}, true}}'
 
-    dtype = f"DataType::dt_{type[0]}"
+    dtype = f"DataType::dt_{DTYPE_NAMES.get(type[0], type[0])}"
     single = len(type) > 1 and type[1] == "single"
     single_str = "true" if single else "false"
     shape = ", ".join(
@@ -255,22 +317,28 @@ def runtime_mixin(commands, device, out_dir):
 
         for name, cmd in commands.items():
             opcode = cmd["opcode"]
+            base_name = cmd["base_name"]
+            suffix = cmd["types_suffix"]
             if cmd.get("custom_op", False):
-                func = f"op_{name}"
+                func = f"op_{base_name}"
+                if "ctype" in cmd:
+                    func += f"<{cmd['ctype']}>"
             else:
                 n_inputs = len(cmd["inputs"])
                 n_outputs = len(cmd["outputs"])
                 dims = cmd.get("dims", 1)
                 vectorized = cmd.get("vectorized", True)
 
+                kernel_name = f"kernel_{base_name}"
                 if device == "cpu":
-                    if vectorized:
-                        kernel = f"kernel_{name}<CpuTypes>, kernel_{name}<SimdTypes>"
-                    else:
-                        kernel = f"kernel_{name}<CpuTypes>, kernel_{name}<CpuTypes>"
+                    simd = "Simd" if vectorized else "Cpu"
+                    kernel = (
+                        f"{kernel_name}<CpuTypes{suffix}>, "
+                        f"{kernel_name}<{simd}Types{suffix}>"
+                    )
                     device_arg = ", DeviceType"
                 elif device == "gpu":
-                    kernel = f"kernel_{name}<GpuTypes>"
+                    kernel = f"{kernel_name}<GpuTypes{suffix}>"
                     device_arg = ""
                 foreach_func = (
                     f"tensor_foreach_dynamic<{kernel}, {n_inputs}, {n_outputs}{device_arg}>"
@@ -293,10 +361,15 @@ def runtime_backward_mixin(commands, device, out_dir):
             if not cmd.get("differentiable", False):
                 continue
             opcode = cmd["opcode"]
+            base_name = cmd["base_name"]
+            suffix = cmd["types_suffix"]
             if cmd.get("custom_op", False):
+                func = f"backward_op_{base_name}"
+                if "ctype" in cmd:
+                    func += f"<{cmd['ctype']}>"
                 f.write(
                     f"case {opcode}:\n"
-                    f"    backward_op_{name}(instr, locals, local_grads, device);\n"
+                    f"    {func}(instr, locals, local_grads, device);\n"
                     f"    break;\n"
                 )
             else:
@@ -318,20 +391,16 @@ def runtime_backward_mixin(commands, device, out_dir):
                 out_stored_str = ",".join(str(i) for i in out_stored)
                 vectorized = cmd.get("vectorized", True)
 
+                kernel_name = f"backward_kernel_{base_name}"
                 if device == "cpu":
-                    if vectorized:
-                        kernel = (
-                            f"backward_kernel_{name}<CpuTypes>, "
-                            f"backward_kernel_{name}<SimdTypes>"
-                        )
-                    else:
-                        kernel = (
-                            f"backward_kernel_{name}<CpuTypes>, "
-                            f"backward_kernel_{name}<CpuTypes>"
-                        )
+                    simd = "Simd" if vectorized else "Cpu"
+                    kernel = (
+                        f"{kernel_name}<CpuTypes{suffix}>, "
+                        f"{kernel_name}<{simd}Types{suffix}>"
+                    )
                     device_arg = ", DeviceType"
                 elif device == "gpu":
-                    kernel = f"backward_kernel_{name}<GpuTypes>"
+                    kernel = f"{kernel_name}<GpuTypes{suffix}>"
                     device_arg = ""
 
                 dims = cmd.get("dims", 1)
@@ -354,8 +423,10 @@ def runtime_backward_mixin(commands, device, out_dir):
 
 
 def is_single_input(arg):
-    return isinstance(arg.get("type"), list) and len(arg["type"]) > 1 and (
-        arg["type"][1] == "single"
+    return (
+        isinstance(arg.get("type"), list)
+        and len(arg["type"]) > 1
+        and (arg["type"][1] == "single")
     )
 
 
@@ -369,7 +440,8 @@ def kernel_checks(commands, out_dir):
         def check(kernel, index, single, what):
             f.write(
                 f"static_assert(\n"
-                f"    single_params<decltype(&{kernel}<SimdTypes>)>::value[{index}] == "
+                f"    single_params<decltype(&{kernel}<SimdTypes{suffix}>)>"
+                f"::value[{index}] == "
                 f"{'true' if single else 'false'},\n"
                 f'    "{kernel}: parameter {index} ({what}) must '
                 f'{"" if single else "not "}be single"\n'
@@ -379,6 +451,8 @@ def check(kernel, index, single, what):
         for name, cmd in commands.items():
             if cmd.get("custom_op", False) or not cmd.get("vectorized", True):
                 continue
+            suffix = cmd["types_suffix"]
+            name = cmd["base_name"]
             inputs = cmd["inputs"]
             if is_single_input(inputs[0]):
                 raise ValueError(
diff --git a/madspace/generate_pyi.py b/madspace/generate_pyi.py
index c1581c1d6b..e2067efb7b 100644
--- a/madspace/generate_pyi.py
+++ b/madspace/generate_pyi.py
@@ -57,6 +57,7 @@ def main():
             "--enum-class-locations=Activation:MLP.Activation",
             "--enum-class-locations=CutMode:Cuts.CutMode",
             "--enum-class-locations=LRSchedule:AdamOptimizer.LRSchedule",
+            "--enum-class-locations=DataType:DataType",
             # Fail the build on unresolvable names/expressions (e.g. an enum
             # default value needing its own --enum-class-locations entry)
             # instead of silently emitting a stub with holes in it.
diff --git a/madspace/include/madspace/compgraphs/function.hpp b/madspace/include/madspace/compgraphs/function.hpp
index 42c0a48dd2..9f373f6d81 100644
--- a/madspace/include/madspace/compgraphs/function.hpp
+++ b/madspace/include/madspace/compgraphs/function.hpp
@@ -128,6 +128,10 @@ class FunctionBuilder {
     ValueVec instruction(const std::string& name, const ValueVec& args);
     /// Append @p instruction with arguments @p args; returns its outputs.
     ValueVec instruction(InstructionPtr instruction, const ValueVec& args);
+    /// Calls the float32 or float64 variant `<name>_float32` / `<name>_float64` of
+    /// an instruction, chosen by the first non-literal float argument. Float64
+    /// literals are converted for a float32 call.
+    ValueVec precision_instruction(const std::string& name, const ValueVec& args);
     /// Index of the execution stream new instructions are added to.
     std::size_t current_stream() const { return _current_stream; }
     /// Select the execution stream new instructions are added to.
diff --git a/madspace/include/madspace/compgraphs/type.hpp b/madspace/include/madspace/compgraphs/type.hpp
index 3084f9dc0d..826b7a47ed 100644
--- a/madspace/include/madspace/compgraphs/type.hpp
+++ b/madspace/include/madspace/compgraphs/type.hpp
@@ -13,20 +13,44 @@ namespace madspace {
 /**
  * Element data type of a compute-graph @ref Value.
  *
- * `dt_int` and `dt_float` are ordinary integer and floating-point tensors.
+ * `dt_int32`, `dt_float32` and `dt_float64` are ordinary integer and
+ * floating-point tensors.
  * `batch_sizes` marks a value that instead carries a list of batch sizes, used
  * to split and recombine batches in multi-channel functions.
  */
 enum class DataType {
-    dt_int,      ///< Integer tensor.
-    dt_float,    ///< Floating-point tensor.
+    dt_int32,    ///< 32-bit integer tensor.
+    dt_float32,  ///< Single-precision floating-point tensor.
+    dt_float64,  ///< Double-precision floating-point tensor.
     batch_sizes, ///< A list of batch sizes rather than tensor data.
 };
 
 using me_int_t = int;
 
 template <typename T>
-concept ScalarType = std::same_as<T, me_int_t> || std::same_as<T, double>;
+concept ScalarType =
+    std::same_as<T, me_int_t> || std::same_as<T, float> || std::same_as<T, double>;
+
+/// Scalar types that can be stored in a literal @ref Value.
+template <typename T>
+concept LiteralType = std::same_as<T, me_int_t> || std::same_as<T, double>;
+
+/// The `DataType` corresponding to the C++ element type `T`.
+template <ScalarType T>
+constexpr DataType dtype_of() {
+    if constexpr (std::is_same_v<T, me_int_t>) {
+        return DataType::dt_int32;
+    } else if constexpr (std::is_same_v<T, float>) {
+        return DataType::dt_float32;
+    } else {
+        return DataType::dt_float64;
+    }
+}
+
+/// True for the floating-point data types `dt_float32` and `dt_float64`.
+inline bool is_float(DataType dtype) {
+    return dtype == DataType::dt_float32 || dtype == DataType::dt_float64;
+}
 
 /**
  * Symbolic size of the batch dimension of a @ref Value.
@@ -138,19 +162,19 @@ inline bool operator!=(const Type& lhs, const Type& rhs) {
 
 using TypeVec = std::vector<Type>;
 
-const Type single_float{DataType::dt_float, BatchSize::One{}, {}};
-const Type single_int{DataType::dt_int, BatchSize::One{}, {}};
+const Type single_float{DataType::dt_float64, BatchSize::One{}, {}};
+const Type single_int{DataType::dt_int32, BatchSize::One{}, {}};
 inline Type single_float_array(int count) {
-    return {DataType::dt_float, BatchSize::one, {count}};
+    return {DataType::dt_float64, BatchSize::one, {count}};
 }
 inline Type single_int_array(int count) {
-    return {DataType::dt_int, BatchSize::one, {count}};
+    return {DataType::dt_int32, BatchSize::one, {count}};
 }
 inline Type single_float_array_2d(int count1, int count2) {
-    return {DataType::dt_float, BatchSize::one, {count1, count2}};
+    return {DataType::dt_float64, BatchSize::one, {count1, count2}};
 }
 inline Type single_int_array_2d(int count1, int count2) {
-    return {DataType::dt_int, BatchSize::one, {count1, count2}};
+    return {DataType::dt_int32, BatchSize::one, {count1, count2}};
 }
 
 const BatchSize batch_size = BatchSize("batch_size");
@@ -158,21 +182,21 @@ Type batch_size_array(int count);
 /// A `batch_sizes` `Type` splitting `batch_size` into `count` symbolic
 /// per-channel sizes that add up to the total.
 Type multichannel_batch_size(int count);
-const Type batch_float{DataType::dt_float, batch_size, {}};
-const Type batch_int{DataType::dt_int, batch_size, {}};
-const Type batch_four_vec{DataType::dt_float, batch_size, {4}};
+const Type batch_float{DataType::dt_float64, batch_size, {}};
+const Type batch_int{DataType::dt_int32, batch_size, {}};
+const Type batch_four_vec{DataType::dt_float64, batch_size, {4}};
 /// A `float` `Type` with a batch dimension and a static trailing shape of
 /// `count` elements.
 inline Type batch_float_array(int count) {
-    return {DataType::dt_float, batch_size, {count}};
+    return {DataType::dt_float64, batch_size, {count}};
 }
 inline Type batch_int_array(int count) {
-    return {DataType::dt_int, batch_size, {count}};
+    return {DataType::dt_int32, batch_size, {count}};
 }
 /// A `Type` of `count` four-vectors: batch dimension plus a trailing shape of
 /// `(count, 4)`.
 inline Type batch_four_vec_array(int count) {
-    return {DataType::dt_float, batch_size, {count, 4}};
+    return {DataType::dt_float64, batch_size, {count, 4}};
 }
 
 using TensorValue = std::tuple<
@@ -207,7 +231,7 @@ struct Value {
     /// Scalar floating-point constant.
     Value(double value) : type(single_float), literal_value(value) {}
 
-    template <ScalarType T>
+    template <LiteralType T>
     Value(const std::vector<std::vector<T>>& values) :
         Value(
             [&] {
@@ -231,10 +255,10 @@ struct Value {
             {static_cast<int>(values.size()), static_cast<int>(values.at(0).size())}
         ) {}
 
-    template <ScalarType T>
+    template <LiteralType T>
     Value(const std::vector<T>& values, const std::vector<int>& shape = {}) :
         type{
-            std::is_same_v<T, me_int_t> ? DataType::dt_int : DataType::dt_float,
+            dtype_of<T>(),
             BatchSize::one,
             shape.size() == 0 ? std::vector<int>{static_cast<int>(values.size())}
                               : shape
diff --git a/madspace/include/madspace/driver/adam_optimizer.hpp b/madspace/include/madspace/driver/adam_optimizer.hpp
index d276f76610..4e78d61441 100644
--- a/madspace/include/madspace/driver/adam_optimizer.hpp
+++ b/madspace/include/madspace/driver/adam_optimizer.hpp
@@ -11,7 +11,8 @@ namespace madspace {
 /// `grad_clip_threshold`.
 class GradientClipper : public FunctionGenerator {
 public:
-    GradientClipper();
+    /// @param dtype  Float type of the gradients; the norm is computed in float64.
+    GradientClipper(DataType dtype = DataType::dt_float64);
 
 private:
     NamedVector<Value> build_function_impl(
diff --git a/madspace/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp
index 9fbad34f1a..408a6c24b6 100644
--- a/madspace/include/madspace/driver/tensor.hpp
+++ b/madspace/include/madspace/driver/tensor.hpp
@@ -446,11 +446,7 @@ class Tensor {
     /// Single-value tensor holding `value`, allocated on `device`.
     template <ScalarType T>
     Tensor(T value, DevicePtr device, AllocHint hint = AllocHint::normal) :
-        impl(new TensorImpl{
-            std::is_same_v<T, me_int_t> ? DataType::dt_int : DataType::dt_float,
-            {1},
-            device
-        }) {
+        impl(new TensorImpl{dtype_of<T>(), {1}, device}) {
         auto size = init_stride();
         allocate(size, *device, hint);
         device->memcpy(impl->data, &value, sizeof(value));
@@ -461,11 +457,7 @@ class Tensor {
 
     template <ScalarType T, typename D>
     Tensor(T value, const D& device, AllocHint hint = AllocHint::normal) :
-        impl(new TensorImpl{
-            std::is_same_v<T, me_int_t> ? DataType::dt_int : DataType::dt_float,
-            {1},
-            device.device_ptr()
-        }) {
+        impl(new TensorImpl{dtype_of<T>(), {1}, device.device_ptr()}) {
         auto size = init_stride();
         allocate(size, device, hint);
         device.memcpy(impl->data, &value, sizeof(value));
@@ -479,8 +471,8 @@ class Tensor {
         impl(new TensorImpl{
             std::visit(
                 Overloaded{
-                    [](std::vector<me_int_t>) { return DataType::dt_int; },
-                    [](std::vector<double>) { return DataType::dt_float; },
+                    [](std::vector<me_int_t>) { return DataType::dt_int32; },
+                    [](std::vector<double>) { return DataType::dt_float64; },
                 },
                 std::get<1>(value)
             ),
@@ -530,6 +522,7 @@ class Tensor {
     TensorView<T, dim> view() {
         check_impl();
         check_view_dim(dim, impl->shape.size());
+        check_view_dtype<T>();
         T* data = static_cast<T*>(impl->data);
         return TensorView<T, dim>(data, impl->stride.data(), impl->shape.data());
     }
@@ -539,6 +532,7 @@ class Tensor {
     const TensorView<T, dim> view() const {
         check_impl();
         check_view_dim(dim, impl->shape.size());
+        check_view_dtype<T>();
         T* data = static_cast<T*>(impl->data);
         return TensorView<T, dim>(data, impl->stride.data(), impl->shape.data());
     }
@@ -547,6 +541,7 @@ class Tensor {
     /// Like @ref view, with the leading `flatten_count` dimensions merged into one.
     PackedTensorView<T, dim> flat_view(std::size_t flatten_count) const {
         check_impl();
+        check_view_dtype<T>();
         T* data = static_cast<T*>(impl->data);
         if (flatten_count <= 1) {
             check_view_dim(dim, impl->shape.size());
@@ -632,9 +627,11 @@ class Tensor {
     std::size_t dtype_size() const {
         check_impl();
         switch (impl->dtype) {
-        case DataType::dt_int:
+        case DataType::dt_int32:
             return sizeof(me_int_t);
-        case DataType::dt_float:
+        case DataType::dt_float32:
+            return sizeof(float);
+        case DataType::dt_float64:
             return sizeof(double);
         case DataType::batch_sizes:
             return 0;
@@ -703,6 +700,9 @@ class Tensor {
     }
     /// A copy of this tensor on the CPU, or `*this` if it is already there.
     Tensor cpu() const { return cpu(*impl->device); }
+    /// A contiguous CPU copy with the element type converted to `dtype`, or the
+    /// CPU tensor itself if it already has that type.
+    Tensor cast(DataType dtype) const;
 
     template <typename D>
     /// Sets every element to zero.
@@ -807,7 +807,7 @@ class Tensor {
 
     /// Maximum number of stored elements (value elements times lanes) of a
     /// literal for which @ref simd_padded_literal pads the storage.
-    static constexpr std::size_t max_simd_padded_size = 256;
+    static constexpr std::size_t max_simd_padded_size = 512;
 
     /**
      * Literal constant for the CPU SIMD kernels.
@@ -921,6 +921,16 @@ class Tensor {
         }
     }
 
+    // reading a tensor through the wrong element type gives garbage without an error
+    template <typename T>
+    void check_view_dtype() const {
+        if constexpr (ScalarType<T>) {
+            if (impl->dtype != dtype_of<T>() && impl->dtype != DataType::batch_sizes) {
+                throw std::invalid_argument("view element type does not match dtype");
+            }
+        }
+    }
+
     template <typename D>
     void allocate(std::size_t size, const D& device, AllocHint hint) {
         auto [data, parent] = device.allocate(size, hint);
diff --git a/madspace/include/madspace/phasespace/channel_weight_network.hpp b/madspace/include/madspace/phasespace/channel_weight_network.hpp
index 23b14cafda..71c597709d 100644
--- a/madspace/include/madspace/phasespace/channel_weight_network.hpp
+++ b/madspace/include/madspace/phasespace/channel_weight_network.hpp
@@ -77,6 +77,7 @@ class ChannelWeightNetwork : public FunctionGenerator {
      * @param prefix                  Prefix for the trainable global names.
      * @param include_preprocessing   If true, apply @ref MomentumPreprocessing
      *                                to the input first.
+     * @param precision               Float type of the @ref MLP.
      */
     ChannelWeightNetwork(
         std::size_t channel_count,
@@ -85,7 +86,8 @@ class ChannelWeightNetwork : public FunctionGenerator {
         std::size_t layers = 3,
         MLP::Activation activation = MLP::leaky_relu,
         const std::string& prefix = "",
-        bool include_preprocessing = true
+        bool include_preprocessing = true,
+        DataType precision = DataType::dt_float64
     );
 
     /// The underlying network.
diff --git a/madspace/include/madspace/phasespace/discrete_flow.hpp b/madspace/include/madspace/phasespace/discrete_flow.hpp
index a4affa9bde..ec990ede60 100644
--- a/madspace/include/madspace/phasespace/discrete_flow.hpp
+++ b/madspace/include/madspace/phasespace/discrete_flow.hpp
@@ -45,6 +45,7 @@ class DiscreteFlow : public Mapping {
      * @param subnet_hidden_dim Hidden width of the @ref MLP subnetworks.
      * @param subnet_layers     Number of layers in the subnetworks.
      * @param subnet_activation  Activation of the subnetworks.
+     * @param subnet_precision  Float type of the subnetworks (see @ref MLP).
      */
     DiscreteFlow(
         const std::vector<std::size_t>& option_counts,
@@ -53,7 +54,8 @@ class DiscreteFlow : public Mapping {
         std::size_t condition_dim = 0,
         std::size_t subnet_hidden_dim = 32,
         std::size_t subnet_layers = 3,
-        MLP::Activation subnet_activation = MLP::leaky_relu
+        MLP::Activation subnet_activation = MLP::leaky_relu,
+        DataType subnet_precision = DataType::dt_float64
     );
     /// Number of categories per discrete dimension.
     const std::vector<std::size_t>& option_counts() const { return _option_counts; }
diff --git a/madspace/include/madspace/phasespace/flow.hpp b/madspace/include/madspace/phasespace/flow.hpp
index 94eeaa3f82..4b267896a2 100644
--- a/madspace/include/madspace/phasespace/flow.hpp
+++ b/madspace/include/madspace/phasespace/flow.hpp
@@ -54,6 +54,7 @@ class Flow : public Mapping {
      * @param subnet_activation  Activation of the subnetworks.
      * @param invert_spline     Whether the spline is applied in the inverse
      *                          direction.
+     * @param subnet_precision  Float type of the subnetworks (see @ref MLP).
      */
     Flow(
         std::size_t input_dim,
@@ -63,7 +64,8 @@ class Flow : public Mapping {
         std::size_t subnet_hidden_dim = 32,
         std::size_t subnet_layers = 3,
         MLP::Activation subnet_activation = MLP::leaky_relu,
-        bool invert_spline = true
+        bool invert_spline = true,
+        DataType subnet_precision = DataType::dt_float64
     );
     /// Number of transformed dimensions.
     std::size_t input_dim() const { return _input_dim; }
diff --git a/madspace/include/madspace/phasespace/mlp.hpp b/madspace/include/madspace/phasespace/mlp.hpp
index 75d34f0bcb..6a63f77f16 100644
--- a/madspace/include/madspace/phasespace/mlp.hpp
+++ b/madspace/include/madspace/phasespace/mlp.hpp
@@ -15,7 +15,9 @@ namespace madspace {
  * between them and a linear output layer. Used as the subnetwork of the
  * @ref Flow and @ref DiscreteFlow adaptive samplers [1]. The weights and biases
  * are compute-graph globals named with @p prefix. Call
- * `initialize_globals(context)` once before use.
+ * `initialize_globals(context)` once before use. With a single-precision
+ * @p precision, the weights and the network evaluation are in float32, while the
+ * input and output stay in float64.
  *
  * `batch` is the leading batch dimension.
  *
@@ -40,18 +42,23 @@ class MLP : public FunctionGenerator {
      * @param layers      Number of layers.
      * @param activation  Activation between the hidden layers.
      * @param prefix      Prefix for the trainable global names.
+     * @param precision   Float type of the weights and the evaluation,
+     *                    `dt_float32` or `dt_float64`.
      */
     MLP(std::size_t input_dim,
         std::size_t output_dim,
         std::size_t hidden_dim = 32,
         std::size_t layers = 3,
         Activation activation = leaky_relu,
-        const std::string& prefix = "");
+        const std::string& prefix = "",
+        DataType precision = DataType::dt_float64);
 
     /// Input width.
     std::size_t input_dim() const { return _input_dim; }
     /// Output width.
     std::size_t output_dim() const { return _output_dim; }
+    /// Float type of the weights and the evaluation.
+    DataType precision() const { return _precision; }
     /**
      * Register the network's trainable parameters on @p context.
      * @param context  Context receiving the globals.
@@ -79,6 +86,7 @@ class MLP : public FunctionGenerator {
     std::size_t _layers;
     Activation _activation;
     std::string _prefix;
+    DataType _precision;
 };
 
 } // namespace madspace
diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml
index 7621b2ff91..cca6c6b3e6 100644
--- a/madspace/instruction_set.yaml
+++ b/madspace/instruction_set.yaml
@@ -231,6 +231,7 @@ sub:
   dims: 0
 
 mul:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -2869,7 +2870,36 @@ interpolate_alpha_s:
 ---
 title: Normalizing flows and MLPs
 
+convert_float32:
+  inputs:
+    - name: in1
+      type: [float64, ...]
+      desc: "double-precision tensor"
+  outputs:
+    - name: out
+      type: [float32, ...]
+      desc: "the input rounded to single precision"
+  dims: 0
+  vectorized: False
+  differentiable: True
+  desc: "Converts a tensor to single precision. Used at the boundaries of single-precision @ref MLP networks."
+
+convert_float64:
+  inputs:
+    - name: in1
+      type: [float32, ...]
+      desc: "single-precision tensor"
+  outputs:
+    - name: out
+      type: [float64, ...]
+      desc: "the input in double precision"
+  dims: 0
+  vectorized: False
+  differentiable: True
+  desc: "Converts a tensor to double precision. Used at the boundaries of single-precision @ref MLP networks."
+
 matmul:
+  precision: [float32, float64]
   inputs:
     - name: x
       type: [float, n]
@@ -2891,6 +2921,7 @@ matmul:
   desc: "Affine layer y = weight x + bias. Building block of @ref MLP."
 
 relu:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -2904,6 +2935,7 @@ relu:
   desc: "Element-wise ReLU activation. Building block of @ref MLP."
 
 leaky_relu:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -2917,6 +2949,7 @@ leaky_relu:
   desc: "Element-wise leaky ReLU activation. Building block of @ref MLP."
 
 elu:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -2930,6 +2963,7 @@ elu:
   desc: "Element-wise ELU activation. Building block of @ref MLP."
 
 gelu:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -2943,6 +2977,7 @@ gelu:
   desc: "Element-wise GELU activation. Building block of @ref MLP."
 
 sigmoid:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -2957,6 +2992,7 @@ sigmoid:
   desc: "Element-wise sigmoid activation. Building block of @ref MLP."
 
 scaled_tanh:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -2971,6 +3007,7 @@ scaled_tanh:
   desc: "Element-wise scaled hyperbolic tangent activation. Building block of @ref MLP."
 
 softplus:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
diff --git a/madspace/src/compgraphs/function.cpp b/madspace/src/compgraphs/function.cpp
index bc2c307d66..f7e07dc382 100644
--- a/madspace/src/compgraphs/function.cpp
+++ b/madspace/src/compgraphs/function.cpp
@@ -263,6 +263,32 @@ ValueVec FunctionBuilder::instruction(const std::string& name, const ValueVec& a
     return instruction(find_instr->second.get(), args);
 }
 
+ValueVec
+FunctionBuilder::precision_instruction(const std::string& name, const ValueVec& args) {
+    auto is_literal = [](const Value& arg) {
+        return !std::holds_alternative<std::monostate>(arg.literal_value);
+    };
+    DataType dtype = DataType::dt_float64;
+    for (auto& arg : args) {
+        if (!is_literal(arg) && is_float(arg.type.dtype)) {
+            dtype = arg.type.dtype;
+            break;
+        }
+    }
+    if (dtype == DataType::dt_float64) {
+        return instruction(name + "_float64", args);
+    }
+    ValueVec converted_args;
+    for (auto& arg : args) {
+        converted_args.push_back(
+            is_literal(arg) && arg.type.dtype == DataType::dt_float64
+                ? convert_float32(arg)
+                : arg
+        );
+    }
+    return instruction(name + "_float32", converted_args);
+}
+
 ValueVec
 FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) {
     // check argument types and determine if all arguments are constant
@@ -315,8 +341,8 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) {
         return arg.local_index == -1 &&
             std::holds_alternative<me_int_t>(arg.literal_value);
     };
-    bool all_float_literals = args.size() > 0 &&
-        std::all_of(args.begin(), args.end(), is_float_literal);
+    bool all_float_literals =
+        args.size() > 0 && std::all_of(args.begin(), args.end(), is_float_literal);
     bool all_int_literals =
         args.size() > 0 && std::all_of(args.begin(), args.end(), is_int_literal);
     if (all_float_literals) {
@@ -328,7 +354,7 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) {
             return {arg(0) + arg(1)};
         case opcodes::sub:
             return {arg(0) - arg(1)};
-        case opcodes::mul:
+        case opcodes::mul_float64:
             return {arg(0) * arg(1)};
         case opcodes::sqrt:
             return {std::sqrt(arg(0))};
@@ -376,7 +402,7 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) {
                     return {var_arg};
                 }
                 break;
-            case opcodes::mul:
+            case opcodes::mul_float64:
                 if (const_val == 1.) {
                     return {var_arg};
                 }
diff --git a/madspace/src/compgraphs/instruction.cpp b/madspace/src/compgraphs/instruction.cpp
index b05d0486ae..f608bfc0d9 100644
--- a/madspace/src/compgraphs/instruction.cpp
+++ b/madspace/src/compgraphs/instruction.cpp
@@ -138,7 +138,7 @@ me_int_t Instruction::int_literal_arg(
 ) const {
     auto& arg_val = args.at(index);
     auto& arg_type = arg_val.type;
-    if (arg_type.dtype != DataType::dt_int || arg_type.batch_size != BatchSize::one ||
+    if (arg_type.dtype != DataType::dt_int32 || arg_type.batch_size != BatchSize::one ||
         arg_type.shape.size() != 0 ||
         !std::holds_alternative<me_int_t>(arg_val.literal_value)) {
         throw std::invalid_argument(
@@ -488,7 +488,7 @@ TypeVec BatchSplitInstruction::signature(const ValueVec& args) const {
 TypeVec BatchSplitByIndexInstruction::signature(const ValueVec& args) const {
     check_arg_count(args, 2);
     auto& indices_type = args.at(0).type;
-    if (indices_type.dtype != DataType::dt_int || indices_type.shape.size() != 0) {
+    if (indices_type.dtype != DataType::dt_int32 || indices_type.shape.size() != 0) {
         throw std::invalid_argument(
             std::format("{}, argument 1: expected batch of integers", name())
         );
@@ -503,7 +503,7 @@ TypeVec BatchSplitByIndexInstruction::signature(const ValueVec& args) const {
     auto last_batch_size = indices_type.batch_size;
     for (int i = 0; i < size; ++i) {
         auto batch_size = i == size - 1 ? last_batch_size : BatchSize();
-        output_types.push_back({DataType::dt_int, batch_size, {}});
+        output_types.push_back({DataType::dt_int32, batch_size, {}});
         last_batch_size = last_batch_size - batch_size;
     }
     return output_types;
@@ -547,7 +547,8 @@ TypeVec BatchMergeByIndexInstruction::signature(const ValueVec& args) const {
                 )
             );
         }
-        if (indices_type.dtype != DataType::dt_int || indices_type.shape.size() != 0) {
+        if (indices_type.dtype != DataType::dt_int32 ||
+            indices_type.shape.size() != 0) {
             throw std::invalid_argument(
                 std::format(
                     "{}, argument {}: expected batch of integers", name(), i + 2
@@ -679,7 +680,7 @@ TypeVec OffsetIndicesInstruction::signature(const ValueVec& args) const {
         batch_sizes_out.type.batch_size_list.end(),
         BatchSize::zero
     );
-    return {{DataType::dt_int, total_batch_size, {}}};
+    return {{DataType::dt_int32, total_batch_size, {}}};
 }
 
 TypeVec FullInstruction::signature(const ValueVec& args) const {
@@ -749,7 +750,7 @@ TypeVec AcceptNormInstruction::signature(const ValueVec& args) const {
         }
         ++i;
     }
-    return {{DataType::dt_float, BatchSize::one, {}}};
+    return {{DataType::dt_float64, BatchSize::one, {}}};
 }
 
 TypeVec RqsReshapeInstruction::signature(const ValueVec& args) const {
@@ -757,7 +758,7 @@ TypeVec RqsReshapeInstruction::signature(const ValueVec& args) const {
     int bin_count = int_literal_arg(args, 1);
 
     auto& input_type = args.at(0).type;
-    if (input_type.dtype != DataType::dt_float || input_type.shape.size() != 1 ||
+    if (input_type.dtype != DataType::dt_float64 || input_type.shape.size() != 1 ||
         input_type.shape.at(0) % (3 * bin_count + 1) != 0) {
         throw std::invalid_argument(
             std::format(
@@ -769,28 +770,28 @@ TypeVec RqsReshapeInstruction::signature(const ValueVec& args) const {
     int dim = input_type.shape.at(0) / (3 * bin_count + 1);
 
     return {
-        {DataType::dt_float, input_type.batch_size, {dim, bin_count}},
-        {DataType::dt_float, input_type.batch_size, {dim, bin_count}},
-        {DataType::dt_float, input_type.batch_size, {dim, bin_count + 1}},
+        {DataType::dt_float64, input_type.batch_size, {dim, bin_count}},
+        {DataType::dt_float64, input_type.batch_size, {dim, bin_count}},
+        {DataType::dt_float64, input_type.batch_size, {dim, bin_count + 1}},
     };
 }
 
 TypeVec NonzeroInstruction::signature(const ValueVec& args) const {
     check_arg_count(args, 1);
     auto& input_type = args.at(0).type;
-    if (input_type.dtype != DataType::dt_float || input_type.shape.size() != 0) {
+    if (input_type.dtype != DataType::dt_float64 || input_type.shape.size() != 0) {
         throw std::invalid_argument(
             std::format("{}, argument 1: expected batch of floats", name())
         );
     }
-    return {{DataType::dt_int, BatchSize(), {}}};
+    return {{DataType::dt_int32, BatchSize(), {}}};
 }
 
 TypeVec BatchGatherInstruction::signature(const ValueVec& args) const {
     check_arg_count(args, 2);
     auto& indices_type = args.at(0).type;
     auto& values_type = args.at(1).type;
-    if (indices_type.dtype != DataType::dt_int || indices_type.shape.size() != 0) {
+    if (indices_type.dtype != DataType::dt_int32 || indices_type.shape.size() != 0) {
         throw std::invalid_argument(
             std::format("{}, argument 1: expected batch of integers", name())
         );
@@ -808,13 +809,13 @@ TypeVec BatchScatterInstruction::signature(const ValueVec& args) const {
     auto& indices_type = args.at(0).type;
     auto& target_type = args.at(1).type;
     auto& source_type = args.at(2).type;
-    if (indices_type.dtype != DataType::dt_int || indices_type.shape.size() != 0) {
+    if (indices_type.dtype != DataType::dt_int32 || indices_type.shape.size() != 0) {
         throw std::invalid_argument(
             std::format("{}, argument 1: expected batch of integers", name())
         );
     }
-    if (target_type.dtype != DataType::dt_float &&
-        target_type.dtype != DataType::dt_int) {
+    if (target_type.dtype != DataType::dt_float64 &&
+        target_type.dtype != DataType::dt_int32) {
         throw std::invalid_argument(
             std::format("{}, argument 2: expected data type float or int", name())
         );
@@ -841,7 +842,7 @@ TypeVec RandomInstruction::signature(const ValueVec& args) const {
         );
     }
     int count = int_literal_arg(args, 1);
-    return {{DataType::dt_float, batch_size_type.batch_size_list.at(0), {count}}};
+    return {{DataType::dt_float64, batch_size_type.batch_size_list.at(0), {count}}};
 }
 
 TypeVec RandomIntInstruction::signature(const ValueVec& args) const {
@@ -860,19 +861,19 @@ TypeVec RandomIntInstruction::signature(const ValueVec& args) const {
             std::format("{}, argument 2: expected single batch size", name())
         );
     }
-    return {{DataType::dt_int, batch_size_type.batch_size_list.at(0), {}}};
+    return {{DataType::dt_int32, batch_size_type.batch_size_list.at(0), {}}};
 }
 
 TypeVec UnweightInstruction::signature(const ValueVec& args) const {
     check_arg_count(args, 2);
     auto& weights_type = args.at(0).type;
     auto& max_weight_type = args.at(1).type;
-    if (weights_type.dtype != DataType::dt_float || weights_type.shape.size() != 0) {
+    if (weights_type.dtype != DataType::dt_float64 || weights_type.shape.size() != 0) {
         throw std::invalid_argument(
             std::format("{}, argument 1: expected batch of floats", name())
         );
     }
-    if (max_weight_type.dtype != DataType::dt_float ||
+    if (max_weight_type.dtype != DataType::dt_float64 ||
         max_weight_type.batch_size != BatchSize::one ||
         max_weight_type.shape.size() != 0) {
         throw std::invalid_argument(
@@ -882,7 +883,8 @@ TypeVec UnweightInstruction::signature(const ValueVec& args) const {
 
     BatchSize out_batch_size;
     return {
-        {DataType::dt_int, out_batch_size, {}}, {DataType::dt_float, out_batch_size, {}}
+        {DataType::dt_int32, out_batch_size, {}},
+        {DataType::dt_float64, out_batch_size, {}}
     };
 }
 
@@ -925,7 +927,7 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const {
         }
         switch (input_key) {
         case UMAMI_IN_MOMENTA:
-            if (input_type.dtype != DataType::dt_float ||
+            if (input_type.dtype != DataType::dt_float64 ||
                 input_type.shape.size() != 2 || input_type.shape.at(1) != 4) {
                 throw std::invalid_argument(
                     std::format(
@@ -939,7 +941,7 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const {
         case UMAMI_IN_RANDOM_COLOR:
         case UMAMI_IN_RANDOM_HELICITY:
         case UMAMI_IN_RANDOM_DIAGRAM:
-            if (input_type.dtype != DataType::dt_float ||
+            if (input_type.dtype != DataType::dt_float64 ||
                 input_type.shape.size() != 0) {
                 throw std::invalid_argument(
                     std::format(
@@ -952,7 +954,8 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const {
         case UMAMI_IN_HELICITY_INDEX:
         case UMAMI_IN_DIAGRAM_INDEX:
         case UMAMI_IN_CHANNEL_INDEX:
-            if (input_type.dtype != DataType::dt_int || input_type.shape.size() != 0) {
+            if (input_type.dtype != DataType::dt_int32 ||
+                input_type.shape.size() != 0) {
                 throw std::invalid_argument(
                     std::format(
                         "matrix_element, argument {}: expected batch of integers", i + 1
@@ -979,17 +982,17 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const {
         std::vector<int> shape;
         switch (output_key) {
         case UMAMI_OUT_MATRIX_ELEMENT:
-            dtype = DataType::dt_float;
+            dtype = DataType::dt_float64;
             shape = {};
             break;
         case UMAMI_OUT_DIAGRAM_AMP2:
-            dtype = DataType::dt_float;
+            dtype = DataType::dt_float64;
             shape = {output_size};
             break;
         case UMAMI_OUT_COLOR_INDEX:
         case UMAMI_OUT_HELICITY_INDEX:
         case UMAMI_OUT_DIAGRAM_INDEX:
-            dtype = DataType::dt_int;
+            dtype = DataType::dt_int32;
             shape = {};
             break;
         default:
diff --git a/madspace/src/compgraphs/type.cpp b/madspace/src/compgraphs/type.cpp
index c82dd051ef..1762d9811d 100644
--- a/madspace/src/compgraphs/type.cpp
+++ b/madspace/src/compgraphs/type.cpp
@@ -86,11 +86,14 @@ BatchSize BatchSize::add(const BatchSize& other, int factor) const {
 
 std::ostream& madspace::operator<<(std::ostream& out, const DataType& dtype) {
     switch (dtype) {
-    case DataType::dt_float:
-        out << "float";
+    case DataType::dt_float32:
+        out << "float32";
         break;
-    case DataType::dt_int:
-        out << "int";
+    case DataType::dt_float64:
+        out << "float64";
+        break;
+    case DataType::dt_int32:
+        out << "int32";
         break;
     case DataType::batch_sizes:
         out << "batch_sizes";
@@ -237,11 +240,14 @@ void madspace::to_json(json& j, const Value& value) {
 
 void madspace::to_json(json& j, const DataType& dtype) {
     switch (dtype) {
-    case DataType::dt_int:
-        j = "int";
+    case DataType::dt_int32:
+        j = "int32";
+        break;
+    case DataType::dt_float32:
+        j = "float32";
         break;
-    case DataType::dt_float:
-        j = "float";
+    case DataType::dt_float64:
+        j = "float64";
         break;
     case DataType::batch_sizes:
         j = "batch_sizes";
@@ -280,14 +286,14 @@ void madspace::from_json(const json& j, Value& value) {
     auto shape = j.at("shape").get<std::vector<int>>();
     auto j_data = j.at("data");
     switch (dtype) {
-    case DataType::dt_int:
+    case DataType::dt_int32:
         if (shape.size() == 0) {
             value = j_data.get<me_int_t>();
         } else {
             value = Value(j_data.get<std::vector<me_int_t>>(), shape);
         }
         break;
-    case DataType::dt_float:
+    case DataType::dt_float64:
         if (shape.size() == 0) {
             value = unescape_special_values(j_data);
         } else {
@@ -298,6 +304,7 @@ void madspace::from_json(const json& j, Value& value) {
             value = Value(tensor, shape);
         }
         break;
+    case DataType::dt_float32:
     case DataType::batch_sizes:
         throw std::invalid_argument("invalid data type");
     }
@@ -305,10 +312,12 @@ void madspace::from_json(const json& j, Value& value) {
 
 void madspace::from_json(const json& j, DataType& dtype) {
     auto str = j.get<std::string>();
-    if (str == "int") {
-        dtype = DataType::dt_int;
-    } else if (str == "float") {
-        dtype = DataType::dt_float;
+    if (str == "int32") {
+        dtype = DataType::dt_int32;
+    } else if (str == "float32") {
+        dtype = DataType::dt_float32;
+    } else if (str == "float64") {
+        dtype = DataType::dt_float64;
     } else if (str == "batch_sizes") {
         dtype = DataType::batch_sizes;
     } else {
diff --git a/madspace/src/cpu/device.cpp b/madspace/src/cpu/device.cpp
index 649fce51ed..1ee1eb9903 100644
--- a/madspace/src/cpu/device.cpp
+++ b/madspace/src/cpu/device.cpp
@@ -8,54 +8,73 @@ using namespace madspace::kernels;
 
 namespace {
 
+// calls func with the scalar and SIMD kernel policies of a float dtype
+template <typename Func>
+void dispatch_float(DataType dtype, const char* op, Func&& func) {
+    switch (dtype) {
+    case DataType::dt_float64:
+        func.template operator()<CpuTypes, SimdTypes>();
+        break;
+    case DataType::dt_float32:
+        func.template operator()<CpuTypes32, SimdTypes32>();
+        break;
+    default:
+        throw std::runtime_error(std::string("invalid dtype in ") + op);
+    }
+}
+
 template <typename D>
 void tensor_copy_impl(const Tensor& source, Tensor& target, const D& device) {
-    if (source.dtype() == DataType::dt_float && target.dtype() == DataType::dt_float) {
-        tensor_foreach_dynamic_single<
-            kernel_copy<CpuTypes>,
-            kernel_copy<SimdTypes>,
-            1,
-            1>({&source}, {&target}, target.size(0), device);
-    } else if (source.dtype() == DataType::dt_int &&
-               target.dtype() == DataType::dt_int) {
+    if (source.dtype() != target.dtype()) {
+        throw std::runtime_error("invalid dtype in copy");
+    }
+    if (source.dtype() == DataType::dt_int32) {
         tensor_foreach_dynamic_single<
             kernel_copy_int<CpuTypes>,
             kernel_copy_int<SimdTypes>,
             1,
             1>({&source}, {&target}, target.size(0), device);
-    } else {
-        throw std::runtime_error("invalid dtype in copy");
+        return;
     }
+    dispatch_float(source.dtype(), "copy", [&]<typename C, typename S>() {
+        tensor_foreach_dynamic_single<kernel_copy<C>, kernel_copy<S>, 1, 1>(
+            {&source}, {&target}, target.size(0), device
+        );
+    });
 }
 
 template <typename D>
 void tensor_zero_impl(Tensor& tensor, const D& device) {
-    if (tensor.dtype() == DataType::dt_float) {
-        tensor_foreach_dynamic_single<
-            kernel_zero<CpuTypes>,
-            kernel_zero<SimdTypes>,
-            1,
-            1>({&tensor}, {&tensor}, tensor.size(0), device);
-    } else if (tensor.dtype() == DataType::dt_int) {
+    if (tensor.dtype() == DataType::dt_int32) {
         tensor_foreach_dynamic_single<
             kernel_zero_int<CpuTypes>,
             kernel_zero_int<SimdTypes>,
             1,
             1>({&tensor}, {&tensor}, tensor.size(0), device);
-    } else {
-        throw std::runtime_error("invalid dtype in zero");
+        return;
     }
+    dispatch_float(tensor.dtype(), "zero", [&]<typename C, typename S>() {
+        tensor_foreach_dynamic_single<kernel_zero<C>, kernel_zero<S>, 1, 1>(
+            {&tensor}, {&tensor}, tensor.size(0), device
+        );
+    });
 }
 
 template <typename D>
 void tensor_add_impl(const Tensor& source, Tensor& target, const D& device) {
+    if (source.dtype() != target.dtype()) {
+        throw std::runtime_error("invalid dtype in add");
+    }
     // a broadcast target (e.g. the gradient of a literal or global) accumulates
     // the whole source batch
-    tensor_foreach_dynamic_single<
-        kernel_add_inplace<CpuTypes>,
-        kernel_add_inplace<SimdTypes>,
-        1,
-        1>({&source}, {&target}, std::max(source.size(0), target.size(0)), device);
+    std::size_t batch_size = std::max(source.size(0), target.size(0));
+    dispatch_float(source.dtype(), "add", [&]<typename C, typename S>() {
+        tensor_foreach_dynamic_single<
+            kernel_add_inplace<C>,
+            kernel_add_inplace<S>,
+            1,
+            1>({&source}, {&target}, batch_size, device);
+    });
 }
 
 } // namespace
@@ -82,22 +101,21 @@ void CpuDevice::adam_step(
     double bias_corr2_sqrt,
     double weight_decay
 ) const {
-    tensor_foreach_dynamic<
-        kernel_adam_step<CpuTypes>,
-        kernel_adam_step<SimdTypes>,
-        1,
-        3>(
-        {&gradient},
-        {&parameter, &exp_avg, &exp_avg_sq},
-        gradient.size(0),
-        *this,
-        step_size,
-        beta1,
-        beta2,
-        eps,
-        bias_corr2_sqrt,
-        weight_decay
-    );
+    dispatch_float(parameter.dtype(), "adam_step", [&]<typename C, typename S>() {
+        using F = typename C::FVal;
+        tensor_foreach_dynamic<kernel_adam_step<C>, kernel_adam_step<S>, 1, 3>(
+            {&gradient},
+            {&parameter, &exp_avg, &exp_avg_sq},
+            gradient.size(0),
+            *this,
+            static_cast<F>(step_size),
+            static_cast<F>(beta1),
+            static_cast<F>(beta2),
+            static_cast<F>(eps),
+            static_cast<F>(bias_corr2_sqrt),
+            static_cast<F>(weight_decay)
+        );
+    });
 }
 
 void AsyncCpuDevice::tensor_copy(const Tensor& source, Tensor& target) const {
diff --git a/madspace/src/cpu/device.hpp b/madspace/src/cpu/device.hpp
index c6de22e360..65ea999413 100644
--- a/madspace/src/cpu/device.hpp
+++ b/madspace/src/cpu/device.hpp
@@ -105,15 +105,18 @@ class AsyncCpuDevice : public CpuDevice {
             return {1, batch_size};
         }
 
-        std::size_t batch_size_vec = (batch_size + simd_vec_size - 1) / simd_vec_size;
-        std::size_t min_batch_size = 64 / simd_vec_size;
+        // jobs are aligned to the largest vector size, so that only the end of the
+        // batch is left to the scalar kernels for every precision
+        std::size_t batch_size_vec =
+            (batch_size + simd_vec_size_max - 1) / simd_vec_size_max;
+        std::size_t min_batch_size = std::max(64 / simd_vec_size_max, 1);
         std::size_t thread_count = _thread_pool.thread_count();
         // std::size_t job_count = (batch_size + min_batch_size - 1) / min_batch_size;
         std::size_t job_count = batch_size_vec < thread_count * min_batch_size
             ? (batch_size_vec + min_batch_size - 1) / min_batch_size
             : thread_count;
         std::size_t job_size =
-            (batch_size_vec + job_count - 1) / job_count * simd_vec_size;
+            (batch_size_vec + job_count - 1) / job_count * simd_vec_size_max;
         // correct rounding errors to vector size
         job_count = (batch_size + job_size - 1) / job_size;
         return {job_count, job_size};
diff --git a/madspace/src/cpu/kernel_definitions.hpp b/madspace/src/cpu/kernel_definitions.hpp
index 79916ef220..bbdc47a5f2 100644
--- a/madspace/src/cpu/kernel_definitions.hpp
+++ b/madspace/src/cpu/kernel_definitions.hpp
@@ -10,32 +10,43 @@
 namespace madspace {
 namespace kernels {
 
-struct CpuTypes {
+template <typename F>
+struct CpuTypesT {
     template <int dim>
-    using FIn = const TensorView<double, dim>;
+    using FIn = const TensorView<F, dim>;
     template <int dim>
     using IIn = const TensorView<me_int_t, dim>;
     template <int dim>
-    using FOut = TensorView<double, dim>;
+    using FOut = TensorView<F, dim>;
     template <int dim>
     using IOut = TensorView<me_int_t, dim>;
     template <int dim>
-    using FSIn = const TensorView<double, dim>;
+    using FSIn = const TensorView<F, dim>;
     template <int dim>
     using ISIn = const TensorView<me_int_t, dim>;
-    using FVal = double;
+    using FVal = F;
     using IVal = me_int_t;
     using BVal = bool;
+    // the same policy with a different float type, for conversion kernels
+    using F32 = CpuTypesT<float>;
+    using F64 = CpuTypesT<double>;
 };
+using CpuTypes = CpuTypesT<double>;
+using CpuTypes32 = CpuTypesT<float>;
 
 inline double where(bool condition, double val_true, double val_false) {
     return condition ? val_true : val_false;
 }
+inline float where(bool condition, float val_true, float val_false) {
+    return condition ? val_true : val_false;
+}
 inline me_int_t where(bool condition, me_int_t val_true, me_int_t val_false) {
     return condition ? val_true : val_false;
 }
 inline double min(double arg1, double arg2) { return arg1 < arg2 ? arg1 : arg2; }
 inline double max(double arg1, double arg2) { return arg1 > arg2 ? arg1 : arg2; }
+inline float min(float arg1, float arg2) { return arg1 < arg2 ? arg1 : arg2; }
+inline float max(float arg1, float arg2) { return arg1 > arg2 ? arg1 : arg2; }
 inline std::size_t single_index(me_int_t arg) { return arg; }
 
 using std::acos;
@@ -79,8 +90,27 @@ struct SimdTypes {
     using IVal = IVec;
     using BVal = BVec;
 };
+
+struct SimdTypes32 {
+    template <int dim>
+    using FIn = const cpu::VectorizedTensorView<FVec32, float, dim, false>;
+    template <int dim>
+    using IIn = const cpu::VectorizedTensorView<IVec32, me_int_t, dim, false>;
+    template <int dim>
+    using FOut = cpu::VectorizedTensorView<FVec32, float, dim, false>;
+    template <int dim>
+    using IOut = cpu::VectorizedTensorView<IVec32, me_int_t, dim, false>;
+    template <int dim>
+    using FSIn = const cpu::SingleVectorizedTensorView<FVec32, float, dim, false>;
+    template <int dim>
+    using ISIn = const cpu::SingleVectorizedTensorView<IVec32, me_int_t, dim, false>;
+    using FVal = FVec32;
+    using IVal = IVec32;
+    using BVal = BVec32;
+};
 #else  // USE_SIMD
 using SimdTypes = CpuTypes;
+using SimdTypes32 = CpuTypes32;
 #endif // USE_SIMD
 
 } // namespace kernels
diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp
index 5685f7d5f5..897f13777e 100644
--- a/madspace/src/cpu/runtime.cpp
+++ b/madspace/src/cpu/runtime.cpp
@@ -29,6 +29,22 @@ extern "C" void dgemm_(
     int* ldc
 );
 
+extern "C" void sgemm_(
+    char* transa,
+    char* transb,
+    int* m,
+    int* n,
+    int* k,
+    float* alpha,
+    float* a,
+    int* lda,
+    float* b,
+    int* ldb,
+    float* beta,
+    float* c,
+    int* ldc
+);
+
 extern "C" void dgemv_(
     char* trans,
     int* m,
@@ -43,6 +59,68 @@ extern "C" void dgemv_(
     int* incy
 );
 
+extern "C" void sgemv_(
+    char* trans,
+    int* m,
+    int* n,
+    float* alpha,
+    float* a,
+    int* lda,
+    float* x,
+    int* incx,
+    float* beta,
+    float* y,
+    int* incy
+);
+
+namespace {
+
+template <typename F>
+void gemm(
+    char transa,
+    char transb,
+    int m,
+    int n,
+    int k,
+    F alpha,
+    F* a,
+    int lda,
+    F* b,
+    int ldb,
+    F beta,
+    F* c,
+    int ldc
+) {
+    if constexpr (std::is_same_v<F, float>) {
+        sgemm_(&transa, &transb, &m, &n, &k, &alpha, a, &lda, b, &ldb, &beta, c, &ldc);
+    } else {
+        dgemm_(&transa, &transb, &m, &n, &k, &alpha, a, &lda, b, &ldb, &beta, c, &ldc);
+    }
+}
+
+template <typename F>
+void gemv(
+    char trans,
+    int m,
+    int n,
+    F alpha,
+    F* a,
+    int lda,
+    F* x,
+    int incx,
+    F beta,
+    F* y,
+    int incy
+) {
+    if constexpr (std::is_same_v<F, float>) {
+        sgemv_(&trans, &m, &n, &alpha, a, &lda, x, &incx, &beta, y, &incy);
+    } else {
+        dgemv_(&trans, &m, &n, &alpha, a, &lda, x, &incx, &beta, y, &incy);
+    }
+}
+
+} // namespace
+
 using namespace madspace;
 using namespace madspace::cpu;
 using namespace madspace::kernels;
@@ -170,7 +248,7 @@ void op_matrix_element(
     );
 }
 
-template <typename D>
+template <typename F, typename D>
 void op_matmul(
     const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device
 ) {
@@ -182,7 +260,7 @@ void op_matmul(
     std::size_t batch_size = input.size(0);
     std::size_t dims_in = input.size(1);
     std::size_t dims_out = weight.size(1);
-    output = Tensor(DataType::dt_float, {batch_size, dims_out}, device);
+    output = Tensor(dtype_of<F>(), {batch_size, dims_out}, device);
     output.copy_from(bias_orig, device);
     if (batch_size == 0) {
         return;
@@ -190,18 +268,25 @@ void op_matmul(
 
     device.sync_barrier();
     device.submit([=]() mutable {
-        char transa = 'N', transb = 'T';
-        int m = batch_size, n = dims_out, k = dims_in;
-        double alpha = 1., beta = 1.;
-        int lda = batch_size, ldb = dims_out, ldc = batch_size;
-        double* a = static_cast<double*>(input.data());
-        double* b = static_cast<double*>(weight.data());
-        double* c = static_cast<double*>(output.data());
-        dgemm_(&transa, &transb, &m, &n, &k, &alpha, a, &lda, b, &ldb, &beta, c, &ldc);
+        gemm<F>(
+            'N',
+            'T',
+            batch_size,
+            dims_out,
+            dims_in,
+            1.,
+            static_cast<F*>(input.data()),
+            batch_size,
+            static_cast<F*>(weight.data()),
+            dims_out,
+            1.,
+            static_cast<F*>(output.data()),
+            batch_size
+        );
     });
 }
 
-template <typename D>
+template <typename F, typename D>
 void backward_op_matmul(
     const CpuRuntime::Instruction& instruction,
     TensorVec& locals,
@@ -220,18 +305,14 @@ void backward_op_matmul(
 
     if (!input_grad) {
         input_grad =
-            Tensor(DataType::dt_float, input.shape(), device, AllocHint::local_grad);
-        // input_grad.zero(device);
+            Tensor(dtype_of<F>(), input.shape(), device, AllocHint::local_grad);
     }
     if (!weight_grad) {
         weight_grad =
-            Tensor(DataType::dt_float, weight.shape(), device, AllocHint::local_grad);
-        // weight_grad.zero(device);
+            Tensor(dtype_of<F>(), weight.shape(), device, AllocHint::local_grad);
     }
     if (!bias_grad) {
-        bias_grad =
-            Tensor(DataType::dt_float, {1, dims_out}, device, AllocHint::local_grad);
-        // bias_grad.zero(device);
+        bias_grad = Tensor(dtype_of<F>(), {1, dims_out}, device, AllocHint::local_grad);
     }
     if (batch_size == 0) {
         return;
@@ -240,42 +321,59 @@ void backward_op_matmul(
 
     // compute input_grad += output_grad * weight
     device.submit([=]() mutable {
-        char transa = 'N', transb = 'N';
-        int m = batch_size, n = dims_in, k = dims_out;
-        double alpha = 1., beta = 1.;
-        int lda = batch_size, ldb = dims_out, ldc = batch_size;
-        double* a = static_cast<double*>(output_grad.data());
-        ;
-        double* b = static_cast<double*>(weight.data());
-        double* c = static_cast<double*>(input_grad.data());
-        ;
-        dgemm_(&transa, &transb, &m, &n, &k, &alpha, a, &lda, b, &ldb, &beta, c, &ldc);
+        gemm<F>(
+            'N',
+            'N',
+            batch_size,
+            dims_in,
+            dims_out,
+            1.,
+            static_cast<F*>(output_grad.data()),
+            batch_size,
+            static_cast<F*>(weight.data()),
+            dims_out,
+            1.,
+            static_cast<F*>(input_grad.data()),
+            batch_size
+        );
     });
 
     // compute weight_grad += output_grad.T * input
     device.submit([=]() mutable {
-        char transa = 'T', transb = 'N';
-        int m = dims_out, n = dims_in, k = batch_size;
-        double alpha = 1., beta = 1.;
-        int lda = batch_size, ldb = batch_size, ldc = dims_out;
-        double* a = static_cast<double*>(output_grad.data());
-        double* b = static_cast<double*>(input.data());
-        double* c = static_cast<double*>(weight_grad.data());
-        dgemm_(&transa, &transb, &m, &n, &k, &alpha, a, &lda, b, &ldb, &beta, c, &ldc);
+        gemm<F>(
+            'T',
+            'N',
+            dims_out,
+            dims_in,
+            batch_size,
+            1.,
+            static_cast<F*>(output_grad.data()),
+            batch_size,
+            static_cast<F*>(input.data()),
+            batch_size,
+            1.,
+            static_cast<F*>(weight_grad.data()),
+            dims_out
+        );
     });
 
     // compute bias_grad += sum_i output_grad_ij
     device.submit([=]() mutable {
         // TODO: we should probably do this differently...
-        std::vector<double> ones(batch_size, 1.);
-        char trans = 'T';
-        int m = batch_size, n = dims_out;
-        double alpha = 1., beta = 1.;
-        int lda = batch_size, incx = 1, incy = 1;
-        double* a = static_cast<double*>(output_grad.data());
-        double* x = ones.data();
-        double* y = static_cast<double*>(bias_grad.data());
-        dgemv_(&trans, &m, &n, &alpha, a, &lda, x, &incx, &beta, y, &incy);
+        std::vector<F> ones(batch_size, 1.);
+        gemv<F>(
+            'T',
+            batch_size,
+            dims_out,
+            1.,
+            static_cast<F*>(output_grad.data()),
+            batch_size,
+            ones.data(),
+            1,
+            1.,
+            static_cast<F*>(bias_grad.data()),
+            1
+        );
     });
 }
 
@@ -287,7 +385,7 @@ void op_nonzero(
     auto& input = locals[instruction.input_indices[0]];
     auto batch_size = input.size(0);
     auto& output = locals[instruction.output_indices[0]];
-    Tensor output_tmp(DataType::dt_int, {batch_size});
+    Tensor output_tmp(DataType::dt_int32, {batch_size});
     auto input_view_flat = input.flat_view<double, 1>(0);
     auto output_view_flat = output_tmp.flat_view<me_int_t, 1>(0);
     device.submit(
@@ -332,13 +430,13 @@ void batch_gather_impl(
 ) {
     Sizes out_shape = values.shape();
     out_shape[0] = indices.size(0);
-    if (values.dtype() == DataType::dt_float) {
-        selection = Tensor(DataType::dt_float, out_shape, device);
+    if (values.dtype() == DataType::dt_float64) {
+        selection = Tensor(DataType::dt_float64, out_shape, device);
         batch_gather_impl_body<kernel_copy<CpuTypes>, dim, double>(
             indices, values, selection, device
         );
-    } else if (values.dtype() == DataType::dt_int) {
-        selection = Tensor(DataType::dt_int, out_shape, device);
+    } else if (values.dtype() == DataType::dt_int32) {
+        selection = Tensor(DataType::dt_int32, out_shape, device);
         batch_gather_impl_body<kernel_copy_int<CpuTypes>, dim, me_int_t>(
             indices, values, selection, device
         );
@@ -396,11 +494,11 @@ template <int dim, typename D>
 void batch_scatter_impl(
     Tensor& indices, Tensor& source, Tensor& output, const D& device
 ) {
-    if (source.dtype() == DataType::dt_float) {
+    if (source.dtype() == DataType::dt_float64) {
         batch_scatter_impl_body<kernel_copy<CpuTypes>, dim, double>(
             indices, source, output, device
         );
-    } else if (source.dtype() == DataType::dt_int) {
+    } else if (source.dtype() == DataType::dt_int32) {
         batch_scatter_impl_body<kernel_copy_int<CpuTypes>, dim, me_int_t>(
             indices, source, output, device
         );
@@ -462,7 +560,7 @@ void op_batch_split_by_index(
         views.reserve(count);
         for (auto [size, output_index] : zip(sizes, instruction.output_indices)) {
             auto& output = locals[output_index];
-            output = Tensor(DataType::dt_int, {size}, device);
+            output = Tensor(DataType::dt_int32, {size}, device);
             views.push_back(output.view<me_int_t, 1>());
         }
         std::fill(sizes.begin(), sizes.end(), 0);
@@ -509,7 +607,7 @@ void batch_reduce_sum_mean_impl(
     auto& input = locals[instruction.input_indices[0]];
     auto& output = locals[instruction.output_indices[0]];
     std::size_t batch_size = input.size(0);
-    output = Tensor(DataType::dt_float, {keepdim ? batch_size : 1}, device);
+    output = Tensor(DataType::dt_float64, {keepdim ? batch_size : 1}, device);
     device.sync_barrier();
 
     auto input_view_flat = input.flat_view<double, 1>(0);
@@ -553,7 +651,7 @@ void batch_reduce_sum_mean_backward_impl(
     auto& output_grad = local_grads[instruction.output_indices[0]];
     if (!input_grad) {
         input_grad =
-            Tensor(DataType::dt_float, input.shape(), device, AllocHint::local_grad);
+            Tensor(DataType::dt_float64, input.shape(), device, AllocHint::local_grad);
         // input_grad.zero(device);
     }
     device.sync_barrier();
@@ -674,7 +772,7 @@ void op_offset_indices(
     auto& sizes_out = locals[instruction.input_indices[1]].batch_sizes();
     std::size_t total_size = std::accumulate(sizes_out.begin(), sizes_out.end(), 0);
     auto& output = locals[instruction.output_indices[0]];
-    output = Tensor(DataType::dt_int, {total_size}, device);
+    output = Tensor(DataType::dt_int32, {total_size}, device);
     auto flat_view = output.flat_view<me_int_t, 1>(0);
 
     std::size_t sum_offset = 0, sum_out = 0;
@@ -709,7 +807,7 @@ void op_quantile(
         return;
     }
 
-    output = Tensor(DataType::dt_float, {1}, device);
+    output = Tensor(DataType::dt_float64, {1}, device);
     device.foreach (
         input.size(0),
         [&](std::size_t count, std::size_t offset) {
@@ -738,7 +836,7 @@ void op_random(
     auto batch_size = locals[instruction.input_indices[0]].batch_sizes()[0];
     auto& output = locals[instruction.output_indices[0]];
     auto dim = instruction.output_shapes[0][0];
-    output = Tensor(DataType::dt_float, {batch_size, dim}, device);
+    output = Tensor(DataType::dt_float64, {batch_size, dim}, device);
     auto flat_view = output.flat_view<double, 1>(2);
     auto& runtime = instruction.runtime;
     device.foreach (
@@ -760,7 +858,7 @@ void op_random_int(
     auto batch_size = locals[instruction.input_indices[0]].batch_sizes()[0];
     auto max_val = locals[instruction.input_indices[1]].batch_sizes()[0];
     auto& output = locals[instruction.output_indices[0]];
-    output = Tensor(DataType::dt_float, {batch_size}, device);
+    output = Tensor(DataType::dt_float64, {batch_size}, device);
     auto flat_view = output.flat_view<me_int_t, 1>(0);
     auto& runtime = instruction.runtime;
     device.foreach (
@@ -788,8 +886,8 @@ void op_unweight(
     auto& uw_weights = locals[instruction.output_indices[1]];
 
     auto batch_size = weights.size(0);
-    Tensor indices_tmp(DataType::dt_int, {batch_size});
-    Tensor uw_weights_tmp(DataType::dt_float, {batch_size});
+    Tensor indices_tmp(DataType::dt_int32, {batch_size});
+    Tensor uw_weights_tmp(DataType::dt_float64, {batch_size});
 
     auto weights_view_flat = weights.flat_view<double, 1>(0);
     auto max_weight_view_flat = max_weight.flat_view<double, 1>(0);
@@ -846,8 +944,8 @@ void op_histogram(
     Sizes shape(out_shape.size() + 1);
     shape[0] = 1;
     std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1);
-    values = Tensor(DataType::dt_float, shape, device);
-    square_values = Tensor(DataType::dt_float, shape, device);
+    values = Tensor(DataType::dt_float64, shape, device);
+    square_values = Tensor(DataType::dt_float64, shape, device);
     device.sync_barrier();
 
     auto weights_view_flat = weights.flat_view<double, 1>(0);
@@ -915,8 +1013,8 @@ void op_vegas_histogram(
     Sizes shape(out_shape.size() + 1);
     shape[0] = 1;
     std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1);
-    values = Tensor(DataType::dt_float, shape, device);
-    counts = Tensor(DataType::dt_int, shape, device);
+    values = Tensor(DataType::dt_float64, shape, device);
+    counts = Tensor(DataType::dt_int32, shape, device);
     device.sync_barrier();
 
     auto input_view_flat = input.flat_view<double, 2>(0);
@@ -974,8 +1072,8 @@ void op_discrete_histogram(
     Sizes shape(out_shape.size() + 1);
     shape[0] = 1;
     std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1);
-    values = Tensor(DataType::dt_float, shape, device);
-    counts = Tensor(DataType::dt_int, shape, device);
+    values = Tensor(DataType::dt_float64, shape, device);
+    counts = Tensor(DataType::dt_int32, shape, device);
 
     std::size_t batch_size = locals[instruction.batch_size_index].size(0);
 
@@ -1134,7 +1232,7 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur
             auto& local = function.locals().at(in.local_index);
             bool is_literal =
                 !std::holds_alternative<std::monostate>(local.literal_value);
-            bool is_int_intermediate = local.type.dtype == DataType::dt_int &&
+            bool is_int_intermediate = local.type.dtype == DataType::dt_int32 &&
                 static_cast<std::size_t>(in.local_index) >= function.inputs().size() &&
                 !global_locals.contains(in.local_index);
             _instructions.back().discard_input_grad.push_back(
@@ -1197,6 +1295,13 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur
         if (context->global_requires_grad(name)) {
             _requires_grad_init.at(value.local_index) = true;
             _grad_global_indices.push_back(value.local_index);
+            // the gradients of all globals are stored in one contiguous tensor
+            if (_grad_global_total_size > 0 && global.dtype() != _grad_global_dtype) {
+                throw std::invalid_argument(
+                    "all globals that require a gradient must have the same dtype"
+                );
+            }
+            _grad_global_dtype = global.dtype();
             _grad_global_shapes.push_back(global.shape());
             _grad_global_total_size += global.shape().product();
         }
@@ -1209,7 +1314,7 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur
                     // padded so that the vectorized kernels can load it like a
                     // batched tensor, see Tensor::simd_padded_literal
                     _locals_init[local.local_index] = Tensor::simd_padded_literal(
-                        LiteralValue(val), simd_vec_size, cpu_device()
+                        LiteralValue(val), simd_vec_size_max, cpu_device()
                     );
                 },
                 [](std::monostate val) {}
@@ -1383,7 +1488,7 @@ std::pair<TensorVec, TensorVec> CpuRuntime::run_backward_single(
     }
 
     Tensor all_global_grads(
-        DataType::dt_float, {_grad_global_total_size}, AllocHint::global_grad
+        _grad_global_dtype, {_grad_global_total_size}, AllocHint::global_grad
     );
     // all_global_grads.zero();
     TensorVec global_grads = all_global_grads.split_and_reshape(_grad_global_shapes);
@@ -1402,9 +1507,9 @@ std::pair<TensorVec, TensorVec> CpuRuntime::run_backward_single(
         for (auto [output_index, output_dtype] :
              zip(instr.output_indices, instr.output_dtypes)) {
             auto& grad = local_grads[output_index];
-            if (!grad && output_dtype == DataType::dt_float) {
+            if (!grad && is_float(output_dtype)) {
                 grad = Tensor(
-                    DataType::dt_float,
+                    output_dtype,
                     locals[output_index].shape(),
                     device,
                     AllocHint::local_grad
@@ -1574,7 +1679,7 @@ std::pair<TensorVec, TensorVec> CpuRuntime::run_backward_concurrent(
     }
 
     Tensor all_global_grads(
-        DataType::dt_float, {_grad_global_total_size}, AllocHint::global_grad
+        _grad_global_dtype, {_grad_global_total_size}, AllocHint::global_grad
     );
     // all_global_grads.zero();
     TensorVec global_grads = all_global_grads.split_and_reshape(_grad_global_shapes);
@@ -1613,9 +1718,9 @@ std::pair<TensorVec, TensorVec> CpuRuntime::run_backward_concurrent(
                     for (auto [output_index, output_dtype] :
                          zip(instr.output_indices, instr.output_dtypes)) {
                         auto& grad = local_grads[output_index];
-                        if (!grad && output_dtype == DataType::dt_float) {
+                        if (!grad && is_float(output_dtype)) {
                             grad = Tensor(
-                                DataType::dt_float,
+                                output_dtype,
                                 locals[output_index].shape(),
                                 device,
                                 AllocHint::local_grad
diff --git a/madspace/src/cpu/runtime.hpp b/madspace/src/cpu/runtime.hpp
index 8a844f79e9..9ea71f994f 100644
--- a/madspace/src/cpu/runtime.hpp
+++ b/madspace/src/cpu/runtime.hpp
@@ -93,6 +93,7 @@ class CpuRuntime : public Runtime {
     SizeVec _grad_global_indices;
     std::vector<Sizes> _grad_global_shapes;
     std::size_t _grad_global_total_size;
+    DataType _grad_global_dtype = DataType::dt_float64;
     ContextPtr _context;
     ThreadResource<MixMaxRandom>& _rand_gens;
     bool _concurrent;
diff --git a/madspace/src/cpu/simd.hpp b/madspace/src/cpu/simd.hpp
index 7f10c491db..6fae8a4392 100644
--- a/madspace/src/cpu/simd.hpp
+++ b/madspace/src/cpu/simd.hpp
@@ -1,5 +1,10 @@
 #pragma once
 
+// Selects the vector type for loads and gathers, which can not be overloaded on
+// the return type alone (e.g. int loads into the float64 or float32 family)
+template <class V>
+struct VecTag {};
+
 #ifdef USE_SIMD
 
 #ifdef USE_SIMD_NEON
@@ -21,5 +26,9 @@
 #else // USE_SIMD
 
 constexpr int simd_vec_size = 1;
+constexpr int simd_vec_size_f32 = 1;
 
 #endif // USE_SIMD
+
+// the largest lane count, used to align jobs and pad literals
+constexpr int simd_vec_size_max = simd_vec_size_f32;
diff --git a/madspace/src/cpu/simd_arm.hpp b/madspace/src/cpu/simd_arm.hpp
index 279ba6c138..d47e6f88ac 100644
--- a/madspace/src/cpu/simd_arm.hpp
+++ b/madspace/src/cpu/simd_arm.hpp
@@ -3,8 +3,11 @@
 #include <sleef.h>
 
 constexpr int simd_vec_size = 2;
+// float32 vectors have twice the lanes of the float64 ones
+constexpr int simd_vec_size_f32 = 4;
 
 struct FVec {
+    static constexpr int lanes = simd_vec_size;
     FVec() = default;
     FVec(float64x2_t _v) : v(_v) {};
     FVec(double _v) : v(vdupq_n_f64(_v)) {};
@@ -18,6 +21,7 @@ struct FVec {
 };
 
 struct IVec {
+    static constexpr int lanes = simd_vec_size;
     IVec() = default;
     IVec(int64x2_t _v) : v(_v) {};
     IVec(int _v) : v(vdupq_n_s64(_v)) {};
@@ -31,6 +35,7 @@ struct IVec {
 };
 
 struct BVec {
+    static constexpr int lanes = simd_vec_size;
     BVec() = default;
     BVec(uint64x2_t _v) : v(_v) {};
     BVec(bool _v) : v(vdupq_n_u64(_v ? ~0ULL : 0)) {};
@@ -39,7 +44,11 @@ struct BVec {
 };
 
 inline FVec vgather(
-    double* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices
+    VecTag<FVec>,
+    double* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec indices
 ) {
     FVec ret;
     IVec strided_indices = indices * IVec(index_stride);
@@ -51,7 +60,11 @@ inline FVec vgather(
 }
 
 inline IVec vgather(
-    int* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices
+    VecTag<IVec>,
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec indices
 ) {
     IVec ret;
     IVec strided_indices = indices * IVec(index_stride);
@@ -63,14 +76,18 @@ inline IVec vgather(
 }
 
 // contiguous loads and stores of the lanes of one batch element
-inline FVec vload(double* base_ptr) { return vld1q_f64(base_ptr); }
+inline FVec vload(VecTag<FVec>, double* base_ptr) { return vld1q_f64(base_ptr); }
 
-inline IVec vload(int* base_ptr) { return vmovl_s32(vld1_s32(base_ptr)); }
+inline IVec vload(VecTag<IVec>, int* base_ptr) { return vmovl_s32(vld1_s32(base_ptr)); }
 
 // the same value in every lane
-inline FVec vload_broadcast(double* base_ptr) { return vld1q_dup_f64(base_ptr); }
+inline FVec vload_broadcast(VecTag<FVec>, double* base_ptr) {
+    return vld1q_dup_f64(base_ptr);
+}
 
-inline IVec vload_broadcast(int* base_ptr) { return vdupq_n_s64(*base_ptr); }
+inline IVec vload_broadcast(VecTag<IVec>, int* base_ptr) {
+    return vdupq_n_s64(*base_ptr);
+}
 
 inline void vscatter(
     double* base_ptr,
@@ -100,7 +117,9 @@ inline void vscatter(
 
 inline void vstore(double* base_ptr, FVec values) { vst1q_f64(base_ptr, values); }
 
-inline void vstore(int* base_ptr, IVec values) { vst1_s32(base_ptr, vmovn_s64(values)); }
+inline void vstore(int* base_ptr, IVec values) {
+    vst1_s32(base_ptr, vmovn_s64(values));
+}
 
 inline FVec where(BVec arg1, FVec arg2, FVec arg3) {
     return vbslq_f64(arg1, arg2, arg3);
@@ -164,3 +183,187 @@ inline FVec log1p(FVec arg1) { return Sleef_log1pd2_u10(arg1); }
 inline FVec expm1(FVec arg1) { return Sleef_expm1d2_u10(arg1); }
 inline FVec erf(FVec arg1) { return Sleef_erfd2_u10(arg1); }
 inline FVec fma(FVec arg1, FVec arg2, FVec arg3) { return vfmaq_f64(arg3, arg1, arg2); }
+
+// float32 family: four lanes, with int32 indices and masks of the same width
+
+struct IVec32;
+
+struct FVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    FVec32() = default;
+    FVec32(float32x4_t _v) : v(_v) {};
+    FVec32(double _v) : v(vdupq_n_f32(static_cast<float>(_v))) {};
+    explicit FVec32(int32x4_t _v) : v(vcvtq_f32_s32(_v)) {};
+    operator float32x4_t() { return v; }
+    FVec32 operator+=(FVec32 _v) {
+        v = vaddq_f32(v, _v);
+        return v;
+    }
+    float32x4_t v;
+};
+
+struct IVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    IVec32() = default;
+    IVec32(int32x4_t _v) : v(_v) {};
+    IVec32(int _v) : v(vdupq_n_s32(_v)) {};
+    explicit IVec32(float32x4_t _v) : v(vcvtq_s32_f32(_v)) {};
+    operator int32x4_t() { return v; }
+    IVec32 operator+=(IVec32 _v) {
+        v = vaddq_s32(v, _v);
+        return v;
+    }
+    int32x4_t v;
+};
+
+struct BVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    BVec32() = default;
+    BVec32(uint32x4_t _v) : v(_v) {};
+    BVec32(bool _v) : v(vdupq_n_u32(_v ? ~0U : 0)) {};
+    operator uint32x4_t() { return v; }
+    uint32x4_t v;
+};
+
+inline IVec32 operator*(IVec32 arg1, IVec32 arg2) { return vmulq_s32(arg1, arg2); }
+
+inline FVec32 vgather(
+    VecTag<FVec32>,
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    int32x4_t idx = indices * IVec32(index_stride);
+    float values[4];
+    for (int i = 0; i < 4; ++i) {
+        values[i] = base_ptr[idx[i] + batch_stride * i];
+    }
+    return vld1q_f32(values);
+}
+
+inline IVec32 vgather(
+    VecTag<IVec32>,
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    int32x4_t idx = indices * IVec32(index_stride);
+    int values[4];
+    for (int i = 0; i < 4; ++i) {
+        values[i] = base_ptr[idx[i] + batch_stride * i];
+    }
+    return vld1q_s32(values);
+}
+
+inline FVec32 vload(VecTag<FVec32>, float* base_ptr) { return vld1q_f32(base_ptr); }
+
+inline IVec32 vload(VecTag<IVec32>, int* base_ptr) { return vld1q_s32(base_ptr); }
+
+inline FVec32 vload_broadcast(VecTag<FVec32>, float* base_ptr) {
+    return vld1q_dup_f32(base_ptr);
+}
+
+inline IVec32 vload_broadcast(VecTag<IVec32>, int* base_ptr) {
+    return vld1q_dup_s32(base_ptr);
+}
+
+inline void vscatter(
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices,
+    FVec32 values
+) {
+    int32x4_t idx = indices * IVec32(index_stride);
+    float32x4_t vals = values;
+    for (int i = 0; i < 4; ++i) {
+        base_ptr[idx[i] + batch_stride * i] = vals[i];
+    }
+}
+
+inline void vscatter(
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices,
+    IVec32 values
+) {
+    int32x4_t idx = indices * IVec32(index_stride);
+    int32x4_t vals = values;
+    for (int i = 0; i < 4; ++i) {
+        base_ptr[idx[i] + batch_stride * i] = vals[i];
+    }
+}
+
+inline void vstore(float* base_ptr, FVec32 values) { vst1q_f32(base_ptr, values); }
+
+inline void vstore(int* base_ptr, IVec32 values) { vst1q_s32(base_ptr, values); }
+
+inline FVec32 where(BVec32 arg1, FVec32 arg2, FVec32 arg3) {
+    return vbslq_f32(arg1, arg2, arg3);
+}
+inline IVec32 where(BVec32 arg1, IVec32 arg2, IVec32 arg3) {
+    return vbslq_s32(arg1, arg2, arg3);
+}
+inline std::size_t single_index(IVec32 arg) { return vgetq_lane_s32(arg, 0); }
+
+inline BVec32 operator==(FVec32 arg1, FVec32 arg2) { return vceqq_f32(arg1, arg2); }
+inline BVec32 operator!=(FVec32 arg1, FVec32 arg2) {
+    return vmvnq_u32(vceqq_f32(arg1, arg2));
+}
+inline BVec32 operator>(FVec32 arg1, FVec32 arg2) { return vcgtq_f32(arg1, arg2); }
+inline BVec32 operator<(FVec32 arg1, FVec32 arg2) { return vcltq_f32(arg1, arg2); }
+inline BVec32 operator>=(FVec32 arg1, FVec32 arg2) { return vcgeq_f32(arg1, arg2); }
+inline BVec32 operator<=(FVec32 arg1, FVec32 arg2) { return vcleq_f32(arg1, arg2); }
+
+inline BVec32 operator==(IVec32 arg1, IVec32 arg2) { return vceqq_s32(arg1, arg2); }
+inline BVec32 operator!=(IVec32 arg1, IVec32 arg2) {
+    return vmvnq_u32(vceqq_s32(arg1, arg2));
+}
+inline BVec32 operator>(IVec32 arg1, IVec32 arg2) { return vcgtq_s32(arg1, arg2); }
+inline BVec32 operator<(IVec32 arg1, IVec32 arg2) { return vcltq_s32(arg1, arg2); }
+inline BVec32 operator>=(IVec32 arg1, IVec32 arg2) { return vcgeq_s32(arg1, arg2); }
+inline BVec32 operator<=(IVec32 arg1, IVec32 arg2) { return vcleq_s32(arg1, arg2); }
+
+inline BVec32 operator&(BVec32 arg1, BVec32 arg2) { return vandq_u32(arg1, arg2); }
+inline BVec32 operator|(BVec32 arg1, BVec32 arg2) { return vorrq_u32(arg1, arg2); }
+inline BVec32 operator!(BVec32 arg1) { return vmvnq_u32(arg1); }
+
+inline FVec32 min(FVec32 arg1, FVec32 arg2) { return where(arg1 < arg2, arg1, arg2); }
+inline FVec32 max(FVec32 arg1, FVec32 arg2) { return where(arg1 > arg2, arg1, arg2); }
+
+inline FVec32 operator-(FVec32 arg1) { return vnegq_f32(arg1); }
+inline FVec32 operator+(FVec32 arg1, FVec32 arg2) { return vaddq_f32(arg1, arg2); }
+inline FVec32 operator-(FVec32 arg1, FVec32 arg2) { return vsubq_f32(arg1, arg2); }
+inline FVec32 operator*(FVec32 arg1, FVec32 arg2) { return vmulq_f32(arg1, arg2); }
+inline FVec32 operator/(FVec32 arg1, FVec32 arg2) { return vdivq_f32(arg1, arg2); }
+inline IVec32 operator-(IVec32 arg1) { return vnegq_s32(arg1); }
+inline IVec32 operator+(IVec32 arg1, IVec32 arg2) { return vaddq_s32(arg1, arg2); }
+inline IVec32 operator-(IVec32 arg1, IVec32 arg2) { return vsubq_s32(arg1, arg2); }
+
+inline BVec32 isnan(FVec32 arg) { return arg != arg; }
+
+inline FVec32 sqrt(FVec32 arg1) { return Sleef_sqrtf4_u05(arg1); }
+inline FVec32 sin(FVec32 arg1) { return Sleef_sinf4_u10(arg1); }
+inline FVec32 cos(FVec32 arg1) { return Sleef_cosf4_u10(arg1); }
+inline FVec32 sinh(FVec32 arg1) { return Sleef_sinhf4_u10(arg1); }
+inline FVec32 asinh(FVec32 arg1) { return Sleef_asinhf4_u10(arg1); }
+inline FVec32 cosh(FVec32 arg1) { return Sleef_coshf4_u10(arg1); }
+inline FVec32 tanh(FVec32 arg1) { return Sleef_tanhf4_u10(arg1); }
+inline FVec32 atan2(FVec32 arg1, FVec32 arg2) { return Sleef_atan2f4_u10(arg1, arg2); }
+inline FVec32 pow(FVec32 arg1, FVec32 arg2) { return Sleef_powf4_u10(arg1, arg2); }
+inline FVec32 fabs(FVec32 arg1) { return Sleef_fabsf4(arg1); }
+inline FVec32 log(FVec32 arg1) { return Sleef_logf4_u10(arg1); }
+inline FVec32 tan(FVec32 arg1) { return Sleef_tanf4_u10(arg1); }
+inline FVec32 atan(FVec32 arg1) { return Sleef_atanf4_u10(arg1); }
+inline FVec32 acos(FVec32 arg1) { return Sleef_acosf4_u10(arg1); }
+inline FVec32 atanh(FVec32 arg1) { return Sleef_atanhf4_u10(arg1); }
+inline FVec32 exp(FVec32 arg1) { return Sleef_expf4_u10(arg1); }
+inline FVec32 log1p(FVec32 arg1) { return Sleef_log1pf4_u10(arg1); }
+inline FVec32 expm1(FVec32 arg1) { return Sleef_expm1f4_u10(arg1); }
+inline FVec32 erf(FVec32 arg1) { return Sleef_erff4_u10(arg1); }
+inline FVec32 fma(FVec32 arg1, FVec32 arg2, FVec32 arg3) {
+    return vfmaq_f32(arg3, arg1, arg2);
+}
diff --git a/madspace/src/cpu/simd_x86_256.hpp b/madspace/src/cpu/simd_x86_256.hpp
index eb770f1158..9cc0e3a0b7 100644
--- a/madspace/src/cpu/simd_x86_256.hpp
+++ b/madspace/src/cpu/simd_x86_256.hpp
@@ -3,6 +3,8 @@
 #include <sleef.h>
 
 constexpr int simd_vec_size = 4;
+// float32 vectors have twice the lanes of the float64 ones
+constexpr int simd_vec_size_f32 = 8;
 
 inline __m128i truncate_i64_to_i32(__m256i arg) {
     __m256 x = _mm256_castsi256_ps(arg);
@@ -12,6 +14,7 @@ inline __m128i truncate_i64_to_i32(__m256i arg) {
 }
 
 struct FVec {
+    static constexpr int lanes = simd_vec_size;
     FVec() = default;
     FVec(__m256d _v) : v(_v) {};
     FVec(double _v) : v(_mm256_set1_pd(_v)) {};
@@ -25,6 +28,7 @@ struct FVec {
 };
 
 struct IVec {
+    static constexpr int lanes = simd_vec_size;
     IVec() = default;
     IVec(__m256i _v) : v(_v) {};
     IVec(int _v) : v(_mm256_set1_epi64x(_v)) {};
@@ -38,6 +42,7 @@ struct IVec {
 };
 
 struct BVec {
+    static constexpr int lanes = simd_vec_size;
     BVec() = default;
     BVec(__m256d _v) : v(_v) {};
     BVec(__m256i _v) : v(_mm256_castsi256_pd(_v)) {};
@@ -60,7 +65,11 @@ mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) {
 }
 
 inline FVec vgather(
-    double* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices
+    VecTag<FVec>,
+    double* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec indices
 ) {
     return _mm256_i32gather_pd(
         base_ptr, mem_indices(batch_stride, index_stride, indices), 8
@@ -68,7 +77,11 @@ inline FVec vgather(
 }
 
 inline IVec vgather(
-    int* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices
+    VecTag<IVec>,
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec indices
 ) {
     return _mm256_cvtepi32_epi64(_mm_i32gather_epi32(
         base_ptr, mem_indices(batch_stride, index_stride, indices), 4
@@ -76,16 +89,20 @@ inline IVec vgather(
 }
 
 // contiguous loads and stores of the lanes of one batch element
-inline FVec vload(double* base_ptr) { return _mm256_loadu_pd(base_ptr); }
+inline FVec vload(VecTag<FVec>, double* base_ptr) { return _mm256_loadu_pd(base_ptr); }
 
-inline IVec vload(int* base_ptr) {
+inline IVec vload(VecTag<IVec>, int* base_ptr) {
     return _mm256_cvtepi32_epi64(_mm_loadu_si128(reinterpret_cast<__m128i*>(base_ptr)));
 }
 
 // the same value in every lane
-inline FVec vload_broadcast(double* base_ptr) { return _mm256_broadcast_sd(base_ptr); }
+inline FVec vload_broadcast(VecTag<FVec>, double* base_ptr) {
+    return _mm256_broadcast_sd(base_ptr);
+}
 
-inline IVec vload_broadcast(int* base_ptr) { return _mm256_set1_epi64x(*base_ptr); }
+inline IVec vload_broadcast(VecTag<IVec>, int* base_ptr) {
+    return _mm256_set1_epi64x(*base_ptr);
+}
 
 inline void vscatter(
     double* base_ptr,
@@ -130,7 +147,9 @@ inline void vscatter(
     }
 }
 
-inline void vstore(double* base_ptr, FVec values) { _mm256_storeu_pd(base_ptr, values); }
+inline void vstore(double* base_ptr, FVec values) {
+    _mm256_storeu_pd(base_ptr, values);
+}
 
 inline void vstore(int* base_ptr, IVec values) {
     _mm_storeu_si128(reinterpret_cast<__m128i*>(base_ptr), truncate_i64_to_i32(values));
@@ -213,3 +232,238 @@ inline FVec erf(FVec arg1) { return Sleef_erfd4_u10avx2(arg1); }
 inline FVec fma(FVec arg1, FVec arg2, FVec arg3) {
     return _mm256_fmadd_pd(arg1, arg2, arg3);
 }
+
+// float32 family: eight lanes, with int32 indices and masks of the same width
+
+struct FVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    FVec32() = default;
+    FVec32(__m256 _v) : v(_v) {};
+    FVec32(double _v) : v(_mm256_set1_ps(static_cast<float>(_v))) {};
+    explicit FVec32(__m256i _v) : v(_mm256_cvtepi32_ps(_v)) {};
+    operator __m256() { return v; }
+    FVec32 operator+=(FVec32 _v) {
+        v = _mm256_add_ps(v, _v);
+        return v;
+    }
+    __m256 v;
+};
+
+struct IVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    IVec32() = default;
+    IVec32(__m256i _v) : v(_v) {};
+    IVec32(int _v) : v(_mm256_set1_epi32(_v)) {};
+    explicit IVec32(__m256 _v) : v(_mm256_cvttps_epi32(_v)) {};
+    operator __m256i() { return v; }
+    IVec32 operator+=(IVec32 _v) {
+        v = _mm256_add_epi32(v, _v);
+        return v;
+    }
+    __m256i v;
+};
+
+struct BVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    BVec32() = default;
+    BVec32(__m256 _v) : v(_v) {};
+    BVec32(__m256i _v) : v(_mm256_castsi256_ps(_v)) {};
+    BVec32(bool _v) : v(_mm256_castsi256_ps(_mm256_set1_epi32(_v ? -1 : 0))) {};
+    operator __m256() { return v; }
+    operator __m256i() { return _mm256_castps_si256(v); }
+    __m256 v;
+};
+
+inline __m256i
+mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) {
+    return _mm256_add_epi32(
+        _mm256_mullo_epi32(indices, _mm256_set1_epi32(index_stride)),
+        _mm256_mullo_epi32(
+            _mm256_set1_epi32(batch_stride), _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0)
+        )
+    );
+}
+
+inline FVec32 vgather(
+    VecTag<FVec32>,
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    return _mm256_i32gather_ps(
+        base_ptr, mem_indices32(batch_stride, index_stride, indices), 4
+    );
+}
+
+inline IVec32 vgather(
+    VecTag<IVec32>,
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    return _mm256_i32gather_epi32(
+        base_ptr, mem_indices32(batch_stride, index_stride, indices), 4
+    );
+}
+
+inline FVec32 vload(VecTag<FVec32>, float* base_ptr) {
+    return _mm256_loadu_ps(base_ptr);
+}
+
+inline IVec32 vload(VecTag<IVec32>, int* base_ptr) {
+    return _mm256_loadu_si256(reinterpret_cast<__m256i*>(base_ptr));
+}
+
+inline FVec32 vload_broadcast(VecTag<FVec32>, float* base_ptr) {
+    return _mm256_broadcast_ss(base_ptr);
+}
+
+inline IVec32 vload_broadcast(VecTag<IVec32>, int* base_ptr) {
+    return _mm256_set1_epi32(*base_ptr);
+}
+
+inline void vscatter(
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices,
+    FVec32 values
+) {
+    float values_buf[8];
+    int indices_buf[8];
+    _mm256_storeu_ps(values_buf, values);
+    _mm256_storeu_si256(
+        reinterpret_cast<__m256i*>(indices_buf),
+        mem_indices32(batch_stride, index_stride, indices)
+    );
+    for (int i = 0; i < 8; ++i) {
+        base_ptr[indices_buf[i]] = values_buf[i];
+    }
+}
+
+inline void vscatter(
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices,
+    IVec32 values
+) {
+    int values_buf[8];
+    int indices_buf[8];
+    _mm256_storeu_si256(reinterpret_cast<__m256i*>(values_buf), values);
+    _mm256_storeu_si256(
+        reinterpret_cast<__m256i*>(indices_buf),
+        mem_indices32(batch_stride, index_stride, indices)
+    );
+    for (int i = 0; i < 8; ++i) {
+        base_ptr[indices_buf[i]] = values_buf[i];
+    }
+}
+
+inline void vstore(float* base_ptr, FVec32 values) {
+    _mm256_storeu_ps(base_ptr, values);
+}
+
+inline void vstore(int* base_ptr, IVec32 values) {
+    _mm256_storeu_si256(reinterpret_cast<__m256i*>(base_ptr), values);
+}
+
+inline FVec32 where(BVec32 arg1, FVec32 arg2, FVec32 arg3) {
+    return _mm256_blendv_ps(arg3, arg2, arg1);
+}
+inline IVec32 where(BVec32 arg1, IVec32 arg2, IVec32 arg3) {
+    return _mm256_blendv_epi8(arg3, arg2, arg1);
+}
+inline std::size_t single_index(IVec32 arg) {
+    return _mm_cvtsi128_si32(_mm256_castsi256_si128(arg));
+}
+inline FVec32 min(FVec32 arg1, FVec32 arg2) { return _mm256_min_ps(arg1, arg2); }
+inline FVec32 max(FVec32 arg1, FVec32 arg2) { return _mm256_max_ps(arg1, arg2); }
+
+inline BVec32 operator==(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps(arg1, arg2, _CMP_EQ_OQ);
+}
+inline BVec32 operator!=(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps(arg1, arg2, _CMP_NEQ_UQ);
+}
+inline BVec32 operator>(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps(arg1, arg2, _CMP_GT_OQ);
+}
+inline BVec32 operator<(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps(arg1, arg2, _CMP_LT_OQ);
+}
+inline BVec32 operator>=(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps(arg1, arg2, _CMP_GE_OQ);
+}
+inline BVec32 operator<=(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps(arg1, arg2, _CMP_LE_OQ);
+}
+
+inline BVec32 operator&(BVec32 arg1, BVec32 arg2) {
+    return _mm256_and_si256(arg1, arg2);
+}
+inline BVec32 operator|(BVec32 arg1, BVec32 arg2) {
+    return _mm256_or_si256(arg1, arg2);
+}
+inline BVec32 operator!(BVec32 arg1) {
+    return _mm256_xor_si256(arg1, _mm256_cmpeq_epi32(arg1, arg1));
+}
+
+inline BVec32 operator==(IVec32 arg1, IVec32 arg2) {
+    return _mm256_cmpeq_epi32(arg1, arg2);
+}
+inline BVec32 operator!=(IVec32 arg1, IVec32 arg2) { return !(arg1 == arg2); }
+inline BVec32 operator>(IVec32 arg1, IVec32 arg2) {
+    return _mm256_cmpgt_epi32(arg1, arg2);
+}
+inline BVec32 operator>=(IVec32 arg1, IVec32 arg2) {
+    return (arg1 > arg2) | (arg1 == arg2);
+}
+inline BVec32 operator<(IVec32 arg1, IVec32 arg2) { return !(arg1 >= arg2); }
+inline BVec32 operator<=(IVec32 arg1, IVec32 arg2) { return !(arg1 > arg2); }
+
+inline FVec32 operator-(FVec32 arg1) {
+    return _mm256_sub_ps(_mm256_set1_ps(0.f), arg1);
+}
+inline FVec32 operator+(FVec32 arg1, FVec32 arg2) { return _mm256_add_ps(arg1, arg2); }
+inline FVec32 operator-(FVec32 arg1, FVec32 arg2) { return _mm256_sub_ps(arg1, arg2); }
+inline FVec32 operator*(FVec32 arg1, FVec32 arg2) { return _mm256_mul_ps(arg1, arg2); }
+inline FVec32 operator/(FVec32 arg1, FVec32 arg2) { return _mm256_div_ps(arg1, arg2); }
+inline IVec32 operator-(IVec32 arg1) {
+    return _mm256_sub_epi32(_mm256_setzero_si256(), arg1);
+}
+inline IVec32 operator+(IVec32 arg1, IVec32 arg2) {
+    return _mm256_add_epi32(arg1, arg2);
+}
+inline IVec32 operator-(IVec32 arg1, IVec32 arg2) {
+    return _mm256_sub_epi32(arg1, arg2);
+}
+
+inline BVec32 isnan(FVec32 arg) { return arg != arg; }
+
+inline FVec32 sqrt(FVec32 arg1) { return Sleef_sqrtf8_u05avx2(arg1); }
+inline FVec32 sin(FVec32 arg1) { return Sleef_sinf8_u10avx2(arg1); }
+inline FVec32 cos(FVec32 arg1) { return Sleef_cosf8_u10avx2(arg1); }
+inline FVec32 sinh(FVec32 arg1) { return Sleef_sinhf8_u10avx2(arg1); }
+inline FVec32 asinh(FVec32 arg1) { return Sleef_asinhf8_u10avx2(arg1); }
+inline FVec32 cosh(FVec32 arg1) { return Sleef_coshf8_u10avx2(arg1); }
+inline FVec32 tanh(FVec32 arg1) { return Sleef_tanhf8_u10avx2(arg1); }
+inline FVec32 atan2(FVec32 arg1, FVec32 arg2) {
+    return Sleef_atan2f8_u10avx2(arg1, arg2);
+}
+inline FVec32 pow(FVec32 arg1, FVec32 arg2) { return Sleef_powf8_u10avx2(arg1, arg2); }
+inline FVec32 fabs(FVec32 arg1) { return Sleef_fabsf8_avx2(arg1); }
+inline FVec32 log(FVec32 arg1) { return Sleef_logf8_u10avx2(arg1); }
+inline FVec32 tan(FVec32 arg1) { return Sleef_tanf8_u10avx2(arg1); }
+inline FVec32 atan(FVec32 arg1) { return Sleef_atanf8_u10avx2(arg1); }
+inline FVec32 acos(FVec32 arg1) { return Sleef_acosf8_u10avx2(arg1); }
+inline FVec32 atanh(FVec32 arg1) { return Sleef_atanhf8_u10avx2(arg1); }
+inline FVec32 exp(FVec32 arg1) { return Sleef_expf8_u10avx2(arg1); }
+inline FVec32 log1p(FVec32 arg1) { return Sleef_log1pf8_u10avx2(arg1); }
+inline FVec32 expm1(FVec32 arg1) { return Sleef_expm1f8_u10avx2(arg1); }
+inline FVec32 erf(FVec32 arg1) { return Sleef_erff8_u10avx2(arg1); }
+inline FVec32 fma(FVec32 arg1, FVec32 arg2, FVec32 arg3) {
+    return _mm256_fmadd_ps(arg1, arg2, arg3);
+}
diff --git a/madspace/src/cpu/simd_x86_256y.hpp b/madspace/src/cpu/simd_x86_256y.hpp
index bdf193a1c6..4d54bc55d0 100644
--- a/madspace/src/cpu/simd_x86_256y.hpp
+++ b/madspace/src/cpu/simd_x86_256y.hpp
@@ -4,8 +4,11 @@
 #include <sleef.h>
 
 constexpr int simd_vec_size = 4;
+// float32 vectors have twice the lanes of the float64 ones
+constexpr int simd_vec_size_f32 = 8;
 
 struct FVec {
+    static constexpr int lanes = simd_vec_size;
     FVec() = default;
     FVec(__m256d _v) : v(_v) {};
     FVec(double _v) : v(_mm256_set1_pd(_v)) {};
@@ -19,6 +22,7 @@ struct FVec {
 };
 
 struct IVec {
+    static constexpr int lanes = simd_vec_size;
     IVec() = default;
     IVec(__m256i _v) : v(_v) {};
     IVec(int _v) : v(_mm256_set1_epi64x(_v)) {};
@@ -32,6 +36,7 @@ struct IVec {
 };
 
 struct BVec {
+    static constexpr int lanes = simd_vec_size;
     BVec() = default;
     BVec(bool _v) : v(_v ? 0x0F : 0) {};
     BVec(__mmask8 _v) : v(_v) {};
@@ -54,7 +59,11 @@ mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) {
 }
 
 inline FVec vgather(
-    double* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices
+    VecTag<FVec>,
+    double* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec indices
 ) {
     return _mm256_i64gather_pd(
         base_ptr, mem_indices(batch_stride, index_stride, indices), 8
@@ -62,7 +71,11 @@ inline FVec vgather(
 }
 
 inline IVec vgather(
-    int* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices
+    VecTag<IVec>,
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec indices
 ) {
     return _mm256_cvtepi32_epi64(_mm256_i64gather_epi32(
         base_ptr, mem_indices(batch_stride, index_stride, indices), 4
@@ -70,16 +83,20 @@ inline IVec vgather(
 }
 
 // contiguous loads and stores of the lanes of one batch element
-inline FVec vload(double* base_ptr) { return _mm256_loadu_pd(base_ptr); }
+inline FVec vload(VecTag<FVec>, double* base_ptr) { return _mm256_loadu_pd(base_ptr); }
 
-inline IVec vload(int* base_ptr) {
+inline IVec vload(VecTag<IVec>, int* base_ptr) {
     return _mm256_cvtepi32_epi64(_mm_loadu_si128(reinterpret_cast<__m128i*>(base_ptr)));
 }
 
 // the same value in every lane
-inline FVec vload_broadcast(double* base_ptr) { return _mm256_broadcast_sd(base_ptr); }
+inline FVec vload_broadcast(VecTag<FVec>, double* base_ptr) {
+    return _mm256_broadcast_sd(base_ptr);
+}
 
-inline IVec vload_broadcast(int* base_ptr) { return _mm256_set1_epi64x(*base_ptr); }
+inline IVec vload_broadcast(VecTag<IVec>, int* base_ptr) {
+    return _mm256_set1_epi64x(*base_ptr);
+}
 
 inline void vscatter(
     double* base_ptr,
@@ -108,7 +125,9 @@ inline void vscatter(
     );
 }
 
-inline void vstore(double* base_ptr, FVec values) { _mm256_storeu_pd(base_ptr, values); }
+inline void vstore(double* base_ptr, FVec values) {
+    _mm256_storeu_pd(base_ptr, values);
+}
 
 inline void vstore(int* base_ptr, IVec values) {
     _mm_storeu_si128(
@@ -209,3 +228,226 @@ inline FVec erf(FVec arg1) { return Sleef_erfd4_u10avx2(arg1); }
 inline FVec fma(FVec arg1, FVec arg2, FVec arg3) {
     return _mm256_fmadd_pd(arg1, arg2, arg3);
 }
+
+// float32 family: eight lanes, with int32 indices and masks of the same width
+
+struct FVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    FVec32() = default;
+    FVec32(__m256 _v) : v(_v) {};
+    FVec32(double _v) : v(_mm256_set1_ps(static_cast<float>(_v))) {};
+    explicit FVec32(__m256i _v) : v(_mm256_cvtepi32_ps(_v)) {};
+    operator __m256() { return v; }
+    FVec32 operator+=(FVec32 _v) {
+        v = _mm256_add_ps(v, _v);
+        return v;
+    }
+    __m256 v;
+};
+
+struct IVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    IVec32() = default;
+    IVec32(__m256i _v) : v(_v) {};
+    IVec32(int _v) : v(_mm256_set1_epi32(_v)) {};
+    explicit IVec32(__m256 _v) : v(_mm256_cvttps_epi32(_v)) {};
+    operator __m256i() { return v; }
+    IVec32 operator+=(IVec32 _v) {
+        v = _mm256_add_epi32(v, _v);
+        return v;
+    }
+    __m256i v;
+};
+
+struct BVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    BVec32() = default;
+    BVec32(bool _v) : v(_v ? 0xFF : 0) {};
+    BVec32(__mmask8 _v) : v(_v) {};
+    operator __mmask8() { return v; }
+    __mmask8 v;
+};
+
+inline __m256i
+mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) {
+    return _mm256_add_epi32(
+        _mm256_mullo_epi32(indices, _mm256_set1_epi32(index_stride)),
+        _mm256_mullo_epi32(
+            _mm256_set1_epi32(batch_stride), _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0)
+        )
+    );
+}
+
+inline FVec32 vgather(
+    VecTag<FVec32>,
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    return _mm256_i32gather_ps(
+        base_ptr, mem_indices32(batch_stride, index_stride, indices), 4
+    );
+}
+
+inline IVec32 vgather(
+    VecTag<IVec32>,
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    return _mm256_i32gather_epi32(
+        base_ptr, mem_indices32(batch_stride, index_stride, indices), 4
+    );
+}
+
+inline FVec32 vload(VecTag<FVec32>, float* base_ptr) {
+    return _mm256_loadu_ps(base_ptr);
+}
+
+inline IVec32 vload(VecTag<IVec32>, int* base_ptr) {
+    return _mm256_loadu_si256(reinterpret_cast<__m256i*>(base_ptr));
+}
+
+inline FVec32 vload_broadcast(VecTag<FVec32>, float* base_ptr) {
+    return _mm256_broadcast_ss(base_ptr);
+}
+
+inline IVec32 vload_broadcast(VecTag<IVec32>, int* base_ptr) {
+    return _mm256_set1_epi32(*base_ptr);
+}
+
+inline void vscatter(
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices,
+    FVec32 values
+) {
+    _mm256_i32scatter_ps(
+        base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4
+    );
+}
+
+inline void vscatter(
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices,
+    IVec32 values
+) {
+    _mm256_i32scatter_epi32(
+        base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4
+    );
+}
+
+inline void vstore(float* base_ptr, FVec32 values) {
+    _mm256_storeu_ps(base_ptr, values);
+}
+
+inline void vstore(int* base_ptr, IVec32 values) {
+    _mm256_storeu_si256(reinterpret_cast<__m256i*>(base_ptr), values);
+}
+
+inline FVec32 where(BVec32 arg1, FVec32 arg2, FVec32 arg3) {
+    return _mm256_mask_blend_ps(arg1, arg3, arg2);
+}
+inline IVec32 where(BVec32 arg1, IVec32 arg2, IVec32 arg3) {
+    return _mm256_mask_blend_epi32(arg1, arg3, arg2);
+}
+inline std::size_t single_index(IVec32 arg) {
+    return _mm_cvtsi128_si32(_mm256_castsi256_si128(arg));
+}
+inline FVec32 min(FVec32 arg1, FVec32 arg2) { return _mm256_min_ps(arg1, arg2); }
+inline FVec32 max(FVec32 arg1, FVec32 arg2) { return _mm256_max_ps(arg1, arg2); }
+
+inline BVec32 operator==(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps_mask(arg1, arg2, _CMP_EQ_OQ);
+}
+inline BVec32 operator!=(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps_mask(arg1, arg2, _CMP_NEQ_UQ);
+}
+inline BVec32 operator>(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps_mask(arg1, arg2, _CMP_GT_OQ);
+}
+inline BVec32 operator<(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps_mask(arg1, arg2, _CMP_LT_OQ);
+}
+inline BVec32 operator>=(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps_mask(arg1, arg2, _CMP_GE_OQ);
+}
+inline BVec32 operator<=(FVec32 arg1, FVec32 arg2) {
+    return _mm256_cmp_ps_mask(arg1, arg2, _CMP_LE_OQ);
+}
+
+inline BVec32 operator&(BVec32 arg1, BVec32 arg2) {
+    return static_cast<__mmask8>(arg1.v & arg2.v);
+}
+inline BVec32 operator|(BVec32 arg1, BVec32 arg2) {
+    return static_cast<__mmask8>(arg1.v | arg2.v);
+}
+inline BVec32 operator!(BVec32 arg1) { return static_cast<__mmask8>(~arg1.v); }
+
+inline BVec32 operator==(IVec32 arg1, IVec32 arg2) {
+    return _mm256_cmpeq_epi32_mask(arg1, arg2);
+}
+inline BVec32 operator!=(IVec32 arg1, IVec32 arg2) {
+    return _mm256_cmpneq_epi32_mask(arg1, arg2);
+}
+inline BVec32 operator>(IVec32 arg1, IVec32 arg2) {
+    return _mm256_cmpgt_epi32_mask(arg1, arg2);
+}
+inline BVec32 operator>=(IVec32 arg1, IVec32 arg2) {
+    return _mm256_cmpge_epi32_mask(arg1, arg2);
+}
+inline BVec32 operator<(IVec32 arg1, IVec32 arg2) {
+    return _mm256_cmplt_epi32_mask(arg1, arg2);
+}
+inline BVec32 operator<=(IVec32 arg1, IVec32 arg2) {
+    return _mm256_cmple_epi32_mask(arg1, arg2);
+}
+
+inline FVec32 operator-(FVec32 arg1) {
+    return _mm256_sub_ps(_mm256_set1_ps(0.f), arg1);
+}
+inline FVec32 operator+(FVec32 arg1, FVec32 arg2) { return _mm256_add_ps(arg1, arg2); }
+inline FVec32 operator-(FVec32 arg1, FVec32 arg2) { return _mm256_sub_ps(arg1, arg2); }
+inline FVec32 operator*(FVec32 arg1, FVec32 arg2) { return _mm256_mul_ps(arg1, arg2); }
+inline FVec32 operator/(FVec32 arg1, FVec32 arg2) { return _mm256_div_ps(arg1, arg2); }
+inline IVec32 operator-(IVec32 arg1) {
+    return _mm256_sub_epi32(_mm256_setzero_si256(), arg1);
+}
+inline IVec32 operator+(IVec32 arg1, IVec32 arg2) {
+    return _mm256_add_epi32(arg1, arg2);
+}
+inline IVec32 operator-(IVec32 arg1, IVec32 arg2) {
+    return _mm256_sub_epi32(arg1, arg2);
+}
+
+inline BVec32 isnan(FVec32 arg) { return arg != arg; }
+
+inline FVec32 sqrt(FVec32 arg1) { return Sleef_sqrtf8_u05avx2(arg1); }
+inline FVec32 sin(FVec32 arg1) { return Sleef_sinf8_u10avx2(arg1); }
+inline FVec32 cos(FVec32 arg1) { return Sleef_cosf8_u10avx2(arg1); }
+inline FVec32 sinh(FVec32 arg1) { return Sleef_sinhf8_u10avx2(arg1); }
+inline FVec32 asinh(FVec32 arg1) { return Sleef_asinhf8_u10avx2(arg1); }
+inline FVec32 cosh(FVec32 arg1) { return Sleef_coshf8_u10avx2(arg1); }
+inline FVec32 tanh(FVec32 arg1) { return Sleef_tanhf8_u10avx2(arg1); }
+inline FVec32 atan2(FVec32 arg1, FVec32 arg2) {
+    return Sleef_atan2f8_u10avx2(arg1, arg2);
+}
+inline FVec32 pow(FVec32 arg1, FVec32 arg2) { return Sleef_powf8_u10avx2(arg1, arg2); }
+inline FVec32 fabs(FVec32 arg1) { return Sleef_fabsf8_avx2(arg1); }
+inline FVec32 log(FVec32 arg1) { return Sleef_logf8_u10avx2(arg1); }
+inline FVec32 tan(FVec32 arg1) { return Sleef_tanf8_u10avx2(arg1); }
+inline FVec32 atan(FVec32 arg1) { return Sleef_atanf8_u10avx2(arg1); }
+inline FVec32 acos(FVec32 arg1) { return Sleef_acosf8_u10avx2(arg1); }
+inline FVec32 atanh(FVec32 arg1) { return Sleef_atanhf8_u10avx2(arg1); }
+inline FVec32 exp(FVec32 arg1) { return Sleef_expf8_u10avx2(arg1); }
+inline FVec32 log1p(FVec32 arg1) { return Sleef_log1pf8_u10avx2(arg1); }
+inline FVec32 expm1(FVec32 arg1) { return Sleef_expm1f8_u10avx2(arg1); }
+inline FVec32 erf(FVec32 arg1) { return Sleef_erff8_u10avx2(arg1); }
+inline FVec32 fma(FVec32 arg1, FVec32 arg2, FVec32 arg3) {
+    return _mm256_fmadd_ps(arg1, arg2, arg3);
+}
diff --git a/madspace/src/cpu/simd_x86_512.hpp b/madspace/src/cpu/simd_x86_512.hpp
index c143fb69a9..0a80a261dc 100644
--- a/madspace/src/cpu/simd_x86_512.hpp
+++ b/madspace/src/cpu/simd_x86_512.hpp
@@ -3,8 +3,11 @@
 #include <sleef.h>
 
 constexpr int simd_vec_size = 8;
+// float32 vectors have twice the lanes of the float64 ones
+constexpr int simd_vec_size_f32 = 16;
 
 struct FVec {
+    static constexpr int lanes = simd_vec_size;
     FVec() = default;
     FVec(__m512d _v) : v(_v) {};
     FVec(double _v) : v(_mm512_set1_pd(_v)) {};
@@ -18,6 +21,7 @@ struct FVec {
 };
 
 struct IVec {
+    static constexpr int lanes = simd_vec_size;
     IVec() = default;
     IVec(__m512i _v) : v(_v) {};
     IVec(int _v) : v(_mm512_set1_epi64(_v)) {};
@@ -31,6 +35,7 @@ struct IVec {
 };
 
 struct BVec {
+    static constexpr int lanes = simd_vec_size;
     BVec() = default;
     BVec(bool _v) : v(_v ? 0xFF : 0) {};
     BVec(__mmask8 _v) : v(_v) {};
@@ -55,7 +60,11 @@ mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) {
 }
 
 inline FVec vgather(
-    double* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices
+    VecTag<FVec>,
+    double* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec indices
 ) {
     return _mm512_i32gather_pd(
         mem_indices(batch_stride, index_stride, indices), base_ptr, 8
@@ -63,7 +72,11 @@ inline FVec vgather(
 }
 
 inline IVec vgather(
-    int* base_ptr, std::size_t batch_stride, std::size_t index_stride, IVec indices
+    VecTag<IVec>,
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec indices
 ) {
     return _mm512_cvtepi32_epi64(_mm256_i32gather_epi32(
         base_ptr, mem_indices(batch_stride, index_stride, indices), 4
@@ -71,18 +84,22 @@ inline IVec vgather(
 }
 
 // contiguous loads and stores of the lanes of one batch element
-inline FVec vload(double* base_ptr) { return _mm512_loadu_pd(base_ptr); }
+inline FVec vload(VecTag<FVec>, double* base_ptr) { return _mm512_loadu_pd(base_ptr); }
 
-inline IVec vload(int* base_ptr) {
+inline IVec vload(VecTag<IVec>, int* base_ptr) {
     return _mm512_cvtepi32_epi64(
         _mm256_loadu_si256(reinterpret_cast<__m256i*>(base_ptr))
     );
 }
 
 // the same value in every lane
-inline FVec vload_broadcast(double* base_ptr) { return _mm512_set1_pd(*base_ptr); }
+inline FVec vload_broadcast(VecTag<FVec>, double* base_ptr) {
+    return _mm512_set1_pd(*base_ptr);
+}
 
-inline IVec vload_broadcast(int* base_ptr) { return _mm512_set1_epi64(*base_ptr); }
+inline IVec vload_broadcast(VecTag<IVec>, int* base_ptr) {
+    return _mm512_set1_epi64(*base_ptr);
+}
 
 inline void vscatter(
     double* base_ptr,
@@ -119,7 +136,9 @@ inline void vscatter(
     }
 }
 
-inline void vstore(double* base_ptr, FVec values) { _mm512_storeu_pd(base_ptr, values); }
+inline void vstore(double* base_ptr, FVec values) {
+    _mm512_storeu_pd(base_ptr, values);
+}
 
 inline void vstore(int* base_ptr, IVec values) {
     _mm256_storeu_si256(
@@ -220,3 +239,229 @@ inline FVec erf(FVec arg1) { return Sleef_erfd8_u10avx512f(arg1); }
 inline FVec fma(FVec arg1, FVec arg2, FVec arg3) {
     return _mm512_fmadd_pd(arg1, arg2, arg3);
 }
+
+// float32 family: sixteen lanes, with int32 indices and masks of the same width
+
+struct FVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    FVec32() = default;
+    FVec32(__m512 _v) : v(_v) {};
+    FVec32(double _v) : v(_mm512_set1_ps(static_cast<float>(_v))) {};
+    explicit FVec32(__m512i _v) : v(_mm512_cvtepi32_ps(_v)) {};
+    operator __m512() { return v; }
+    FVec32 operator+=(FVec32 _v) {
+        v = _mm512_add_ps(v, _v);
+        return v;
+    }
+    __m512 v;
+};
+
+struct IVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    IVec32() = default;
+    IVec32(__m512i _v) : v(_v) {};
+    IVec32(int _v) : v(_mm512_set1_epi32(_v)) {};
+    explicit IVec32(__m512 _v) : v(_mm512_cvttps_epi32(_v)) {};
+    operator __m512i() { return v; }
+    IVec32 operator+=(IVec32 _v) {
+        v = _mm512_add_epi32(v, _v);
+        return v;
+    }
+    __m512i v;
+};
+
+struct BVec32 {
+    static constexpr int lanes = simd_vec_size_f32;
+    BVec32() = default;
+    BVec32(bool _v) : v(_v ? 0xFFFF : 0) {};
+    BVec32(__mmask16 _v) : v(_v) {};
+    operator __mmask16() { return v; }
+    __mmask16 v;
+};
+
+inline __m512i
+mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) {
+    return _mm512_add_epi32(
+        _mm512_mullo_epi32(indices, _mm512_set1_epi32(index_stride)),
+        _mm512_mullo_epi32(
+            _mm512_set1_epi32(batch_stride),
+            _mm512_set_epi32(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0)
+        )
+    );
+}
+
+inline FVec32 vgather(
+    VecTag<FVec32>,
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    return _mm512_i32gather_ps(
+        mem_indices32(batch_stride, index_stride, indices), base_ptr, 4
+    );
+}
+
+inline IVec32 vgather(
+    VecTag<IVec32>,
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    return _mm512_i32gather_epi32(
+        mem_indices32(batch_stride, index_stride, indices), base_ptr, 4
+    );
+}
+
+inline FVec32 vload(VecTag<FVec32>, float* base_ptr) {
+    return _mm512_loadu_ps(base_ptr);
+}
+
+inline IVec32 vload(VecTag<IVec32>, int* base_ptr) {
+    return _mm512_loadu_si512(base_ptr);
+}
+
+inline FVec32 vload_broadcast(VecTag<FVec32>, float* base_ptr) {
+    return _mm512_set1_ps(*base_ptr);
+}
+
+inline IVec32 vload_broadcast(VecTag<IVec32>, int* base_ptr) {
+    return _mm512_set1_epi32(*base_ptr);
+}
+
+inline void vscatter(
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices,
+    FVec32 values
+) {
+    _mm512_i32scatter_ps(
+        base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4
+    );
+}
+
+inline void vscatter(
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices,
+    IVec32 values
+) {
+    _mm512_i32scatter_epi32(
+        base_ptr, mem_indices32(batch_stride, index_stride, indices), values, 4
+    );
+}
+
+inline void vstore(float* base_ptr, FVec32 values) {
+    _mm512_storeu_ps(base_ptr, values);
+}
+
+inline void vstore(int* base_ptr, IVec32 values) {
+    _mm512_storeu_si512(base_ptr, values);
+}
+
+inline FVec32 where(BVec32 arg1, FVec32 arg2, FVec32 arg3) {
+    return _mm512_mask_blend_ps(arg1, arg3, arg2);
+}
+inline IVec32 where(BVec32 arg1, IVec32 arg2, IVec32 arg3) {
+    return _mm512_mask_blend_epi32(arg1, arg3, arg2);
+}
+inline std::size_t single_index(IVec32 arg) {
+    return _mm_cvtsi128_si32(_mm512_castsi512_si128(arg));
+}
+inline FVec32 min(FVec32 arg1, FVec32 arg2) { return _mm512_min_ps(arg1, arg2); }
+inline FVec32 max(FVec32 arg1, FVec32 arg2) { return _mm512_max_ps(arg1, arg2); }
+
+inline BVec32 operator==(FVec32 arg1, FVec32 arg2) {
+    return _mm512_cmp_ps_mask(arg1, arg2, _CMP_EQ_OQ);
+}
+inline BVec32 operator!=(FVec32 arg1, FVec32 arg2) {
+    return _mm512_cmp_ps_mask(arg1, arg2, _CMP_NEQ_UQ);
+}
+inline BVec32 operator>(FVec32 arg1, FVec32 arg2) {
+    return _mm512_cmp_ps_mask(arg1, arg2, _CMP_GT_OQ);
+}
+inline BVec32 operator<(FVec32 arg1, FVec32 arg2) {
+    return _mm512_cmp_ps_mask(arg1, arg2, _CMP_LT_OQ);
+}
+inline BVec32 operator>=(FVec32 arg1, FVec32 arg2) {
+    return _mm512_cmp_ps_mask(arg1, arg2, _CMP_GE_OQ);
+}
+inline BVec32 operator<=(FVec32 arg1, FVec32 arg2) {
+    return _mm512_cmp_ps_mask(arg1, arg2, _CMP_LE_OQ);
+}
+
+inline BVec32 operator&(BVec32 arg1, BVec32 arg2) {
+    return static_cast<__mmask16>(arg1.v & arg2.v);
+}
+inline BVec32 operator|(BVec32 arg1, BVec32 arg2) {
+    return static_cast<__mmask16>(arg1.v | arg2.v);
+}
+inline BVec32 operator!(BVec32 arg1) { return static_cast<__mmask16>(~arg1.v); }
+
+inline BVec32 operator==(IVec32 arg1, IVec32 arg2) {
+    return _mm512_cmpeq_epi32_mask(arg1, arg2);
+}
+inline BVec32 operator!=(IVec32 arg1, IVec32 arg2) {
+    return _mm512_cmpneq_epi32_mask(arg1, arg2);
+}
+inline BVec32 operator>(IVec32 arg1, IVec32 arg2) {
+    return _mm512_cmpgt_epi32_mask(arg1, arg2);
+}
+inline BVec32 operator>=(IVec32 arg1, IVec32 arg2) {
+    return _mm512_cmpge_epi32_mask(arg1, arg2);
+}
+inline BVec32 operator<(IVec32 arg1, IVec32 arg2) {
+    return _mm512_cmplt_epi32_mask(arg1, arg2);
+}
+inline BVec32 operator<=(IVec32 arg1, IVec32 arg2) {
+    return _mm512_cmple_epi32_mask(arg1, arg2);
+}
+
+inline FVec32 operator-(FVec32 arg1) {
+    return _mm512_sub_ps(_mm512_set1_ps(0.f), arg1);
+}
+inline FVec32 operator+(FVec32 arg1, FVec32 arg2) { return _mm512_add_ps(arg1, arg2); }
+inline FVec32 operator-(FVec32 arg1, FVec32 arg2) { return _mm512_sub_ps(arg1, arg2); }
+inline FVec32 operator*(FVec32 arg1, FVec32 arg2) { return _mm512_mul_ps(arg1, arg2); }
+inline FVec32 operator/(FVec32 arg1, FVec32 arg2) { return _mm512_div_ps(arg1, arg2); }
+inline IVec32 operator-(IVec32 arg1) {
+    return _mm512_sub_epi32(_mm512_setzero_si512(), arg1);
+}
+inline IVec32 operator+(IVec32 arg1, IVec32 arg2) {
+    return _mm512_add_epi32(arg1, arg2);
+}
+inline IVec32 operator-(IVec32 arg1, IVec32 arg2) {
+    return _mm512_sub_epi32(arg1, arg2);
+}
+
+inline BVec32 isnan(FVec32 arg) { return arg != arg; }
+
+inline FVec32 sqrt(FVec32 arg1) { return Sleef_sqrtf16_u05avx512f(arg1); }
+inline FVec32 sin(FVec32 arg1) { return Sleef_sinf16_u10avx512f(arg1); }
+inline FVec32 cos(FVec32 arg1) { return Sleef_cosf16_u10avx512f(arg1); }
+inline FVec32 sinh(FVec32 arg1) { return Sleef_sinhf16_u10avx512f(arg1); }
+inline FVec32 asinh(FVec32 arg1) { return Sleef_asinhf16_u10avx512f(arg1); }
+inline FVec32 cosh(FVec32 arg1) { return Sleef_coshf16_u10avx512f(arg1); }
+inline FVec32 tanh(FVec32 arg1) { return Sleef_tanhf16_u10avx512f(arg1); }
+inline FVec32 atan2(FVec32 arg1, FVec32 arg2) {
+    return Sleef_atan2f16_u10avx512f(arg1, arg2);
+}
+inline FVec32 pow(FVec32 arg1, FVec32 arg2) {
+    return Sleef_powf16_u10avx512f(arg1, arg2);
+}
+inline FVec32 fabs(FVec32 arg1) { return Sleef_fabsf16_avx512f(arg1); }
+inline FVec32 log(FVec32 arg1) { return Sleef_logf16_u10avx512f(arg1); }
+inline FVec32 tan(FVec32 arg1) { return Sleef_tanf16_u10avx512f(arg1); }
+inline FVec32 atan(FVec32 arg1) { return Sleef_atanf16_u10avx512f(arg1); }
+inline FVec32 acos(FVec32 arg1) { return Sleef_acosf16_u10avx512f(arg1); }
+inline FVec32 atanh(FVec32 arg1) { return Sleef_atanhf16_u10avx512f(arg1); }
+inline FVec32 exp(FVec32 arg1) { return Sleef_expf16_u10avx512f(arg1); }
+inline FVec32 log1p(FVec32 arg1) { return Sleef_log1pf16_u10avx512f(arg1); }
+inline FVec32 expm1(FVec32 arg1) { return Sleef_expm1f16_u10avx512f(arg1); }
+inline FVec32 erf(FVec32 arg1) { return Sleef_erff16_u10avx512f(arg1); }
+inline FVec32 fma(FVec32 arg1, FVec32 arg2, FVec32 arg3) {
+    return _mm512_fmadd_ps(arg1, arg2, arg3);
+}
diff --git a/madspace/src/cpu/tensor.hpp b/madspace/src/cpu/tensor.hpp
index 0eeba4ba86..00d8345550 100644
--- a/madspace/src/cpu/tensor.hpp
+++ b/madspace/src/cpu/tensor.hpp
@@ -12,12 +12,12 @@
 namespace madspace {
 namespace cpu {
 
-// View onto a tensor for the vectorized kernels. The simd_vec_size lanes of a
+// View onto a tensor for the vectorized kernels. The V::lanes lanes of a
 // vector are consecutive batch entries, so they are always contiguous in memory
 // (batch stride one) and are read and written with plain vector loads and stores.
 // Moving to the next vector advances the data pointer by vec_step elements:
-// simd_vec_size for batched tensors, zero for literals that are stored padded to
-// simd_vec_size identical lanes. Tensors that do not have this layout are handled
+// V::lanes for batched tensors, zero for literals that are stored padded to
+// simd_vec_size_max identical lanes. Tensors that do not have this layout are handled
 // by the scalar kernels instead (see tensor_foreach_impl).
 template <class V, ScalarType T, int _dim, bool is_batch>
 class VectorizedTensorView {
@@ -65,9 +65,7 @@ class VectorizedTensorView {
         T* ptr = &_data[offset(first_index)];
         int i = 1;
         ((ptr = &ptr[index * _stride[i++]]), ...);
-        return {
-            ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _vec_step
-        };
+        return {ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _vec_step};
     }
 
     template <typename... I>
@@ -78,15 +76,13 @@ class VectorizedTensorView {
         T* ptr = &_data[offset(first_index)];
         int i = 1;
         ((ptr = &ptr[index * _stride[i++]]), ...);
-        return {
-            ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _vec_step
-        };
+        return {ptr, &_stride[sizeof...(I) + 1], &_shape[sizeof...(I) + 1], _vec_step};
     }
 
     operator V() const
         requires(_dim == 0)
     {
-        return vload(_data);
+        return vload(VecTag<V>{}, _data);
     }
 
     V operator=(V value)
@@ -99,7 +95,7 @@ class VectorizedTensorView {
     V operator+=(V value)
         requires(_dim == 0)
     {
-        V new_value = vload(_data) + value;
+        V new_value = vload(VecTag<V>{}, _data) + value;
         vstore(_data, new_value);
         return new_value;
     }
@@ -109,7 +105,7 @@ class VectorizedTensorView {
 
     std::size_t size(std::size_t index = 0) const {
         if (is_batch && index == 0) {
-            return _shape[0] / simd_vec_size;
+            return _shape[0] / V::lanes;
         } else {
             return _shape[index];
         }
@@ -119,14 +115,14 @@ class VectorizedTensorView {
     V gather(IVec indices) const
         requires(_dim == 1)
     {
-        return vgather(_data, 1, _stride[0], indices);
+        return vgather(VecTag<V>{}, _data, 1, _stride[0], indices);
     }
 
     template <typename IVec>
     void scatter_add(IVec indices, V values)
         requires(_dim == 1)
     {
-        V old_values = vgather(_data, 1, _stride[0], indices);
+        V old_values = vgather(VecTag<V>{}, _data, 1, _stride[0], indices);
         vscatter(_data, 1, _stride[0], indices, old_values + values);
     }
 
@@ -183,7 +179,7 @@ class SingleVectorizedTensorView {
     operator V() const
         requires(_dim == 0)
     {
-        return vload_broadcast(_data);
+        return vload_broadcast(VecTag<V>{}, _data);
     }
 
     // the value as a plain scalar, e.g. for loop bounds or indices
@@ -199,7 +195,7 @@ class SingleVectorizedTensorView {
     V gather(IVec indices) const
         requires(_dim == 1)
     {
-        return vgather(_data, 0, _stride[0], indices);
+        return vgather(VecTag<V>{}, _data, 0, _stride[0], indices);
     }
 
 private:
@@ -277,7 +273,7 @@ struct get_views {
 
 // return the tuple of VectorizedTensorViews where the type is extracted
 // from the signature of F. vec_steps holds the step between vectors for every
-// argument (simd_vec_size, or zero for padded literals)
+// argument (the lane count, or zero for padded literals)
 template <typename F, int dims>
 struct get_vectorized_views;
 template <typename... TParam, int dims>
@@ -336,6 +332,15 @@ struct first_param<void (*)(TParam...)> {
     static constexpr int dim = std::tuple_element_t<0, std::tuple<TParam...>>::dim;
 };
 
+// lane count of a vectorized kernel, taken from the vector type of its first view
+template <typename F>
+struct vector_lanes;
+template <typename... TParam>
+struct vector_lanes<void (*)(TParam...)> {
+    static constexpr int value = std::remove_cvref_t<
+        std::tuple_element_t<0, std::tuple<TParam...>>>::VType::lanes;
+};
+
 template <auto func, int dims, typename... V>
 inline void nested_for(std::size_t batch_size, std::size_t batch_offset, V... views) {
     std::size_t end_index = batch_offset + batch_size;
@@ -437,7 +442,7 @@ inline void tensor_foreach_impl(
 
     // The vectorized kernel loads and stores the lanes of a vector with plain
     // contiguous memory accesses. It is only used if every batched argument has
-    // batch stride one (or is a literal padded to simd_vec_size lanes); single
+    // batch stride one (or is a literal padded to simd_vec_size_max lanes); single
     // inputs are read with broadcast loads and can have any layout. Otherwise
     // the whole batch is handled by the scalar kernel.
     constexpr std::size_t n_args = n_in + n_out + sizeof...(S);
@@ -454,7 +459,7 @@ inline void tensor_foreach_impl(
                     if constexpr (!ScalarType<View> && !Single::value[I]) {
                         auto& view = std::get<I>(flat_views);
                         if (view.shape[0] == batch_size && view.stride[0] == 1) {
-                            vec_steps[I] = simd_vec_size;
+                            vec_steps[I] = vector_lanes<decltype(vector_func)>::value;
                         } else if (view.simd_padded) {
                             vec_steps[I] = 0;
                         } else {
@@ -469,10 +474,9 @@ inline void tensor_foreach_impl(
 
     device.foreach (
         batch_size,
-        [flat_views,
-         vectorize,
-         vec_steps,
-         scalar_args...](std::size_t count, std::size_t offset) mutable {
+        [flat_views, vectorize, vec_steps, scalar_args...](
+            std::size_t count, std::size_t offset
+        ) mutable {
             auto views = std::apply(get_views(), flat_views);
             std::size_t scalar_offset = offset;
             if constexpr (!std::
@@ -486,8 +490,10 @@ inline void tensor_foreach_impl(
                         },
                         views
                     );
-                    std::size_t vec_count = count / simd_vec_size;
-                    std::size_t vec_offset = offset / simd_vec_size;
+                    constexpr std::size_t lanes =
+                        vector_lanes<decltype(vector_func)>::value;
+                    std::size_t vec_count = count / lanes;
+                    std::size_t vec_offset = offset / lanes;
                     std::apply(
                         [vec_count, vec_offset](auto&&... args) {
                             nested_for<vector_func, dims>(
@@ -496,7 +502,7 @@ inline void tensor_foreach_impl(
                         },
                         vectorized_views
                     );
-                    scalar_offset += vec_count * simd_vec_size;
+                    scalar_offset += vec_count * lanes;
                 }
             }
             std::size_t scalar_count = offset + count - scalar_offset;
diff --git a/madspace/src/driver/adam_optimizer.cpp b/madspace/src/driver/adam_optimizer.cpp
index 7e05fe7430..f2a9380fff 100644
--- a/madspace/src/driver/adam_optimizer.cpp
+++ b/madspace/src/driver/adam_optimizer.cpp
@@ -4,20 +4,24 @@
 
 using namespace madspace;
 
-GradientClipper::GradientClipper() :
+GradientClipper::GradientClipper(DataType dtype) :
     FunctionGenerator(
         "Unweighter",
-        {{"gradients_in", batch_float}, {"threshold", single_float}},
-        {{"gradients_out", batch_float}, {"gradient_norm", single_float}}
+        {{"gradients_in", Type(dtype, batch_size, {})}, {"threshold", single_float}},
+        {{"gradients_out", Type(dtype, batch_size, {})},
+         {"gradient_norm", single_float}}
     ) {}
 
 NamedVector<Value> GradientClipper::build_function_impl(
     FunctionBuilder& fb, const NamedVector<Value>& args
 ) const {
+    // only the norm is computed in double precision
     Value grads_in = args.at(0);
-    Value grad_norm = fb.sqrt(fb.batch_reduce_sum(fb.square(grads_in)));
+    bool single = grads_in.type.dtype == DataType::dt_float32;
+    Value grads_sq = fb.square(single ? fb.convert_float64(grads_in) : grads_in);
+    Value grad_norm = fb.sqrt(fb.batch_reduce_sum(grads_sq));
     Value factor = fb.min(fb.div(args.at(1), grad_norm), 1.0);
-    Value grads_out = fb.mul(grads_in, factor);
+    Value grads_out = fb.mul(grads_in, single ? fb.convert_float32(factor) : factor);
     return {
         {"gradients_out", grads_out},
         {"gradient_norm", grad_norm},
@@ -46,11 +50,6 @@ AdamOptimizer::AdamOptimizer(
     _eps(eps),
     _grad_clip_threshold(grad_clip_threshold),
     _weight_decay(weight_decay),
-    _grad_clipper(
-        grad_clip_threshold > 0.
-            ? build_runtime(GradientClipper().function(), context, false)
-            : nullptr
-    ),
     _one(1.0, context->device()),
     _loss_mean(std::numeric_limits<double>::quiet_NaN()) {
     DevicePtr device = context->device();
@@ -61,6 +60,11 @@ AdamOptimizer::AdamOptimizer(
     }
     _parameter = context->reallocate_globals_contiguously(_param_names);
     _runtime = build_runtime(function, context);
+    if (grad_clip_threshold > 0.) {
+        _grad_clipper = build_runtime(
+            GradientClipper(_parameter.dtype()).function(), context, false
+        );
+    }
     _exp_avg = Tensor(_parameter.dtype(), _parameter.shape(), _parameter.device());
     _exp_avg.zero();
     _exp_avg_sq = Tensor(_parameter.dtype(), _parameter.shape(), _parameter.device());
@@ -141,6 +145,11 @@ void AdamOptimizer::replace_function(const Function& function) {
     _runtime.reset();
     _parameter = _context->reallocate_globals_contiguously(_param_names);
     _runtime = build_runtime(function, _context);
+    if (_grad_clipper) {
+        _grad_clipper = build_runtime(
+            GradientClipper(_parameter.dtype()).function(), _context, false
+        );
+    }
     Tensor old_ea = _exp_avg, old_eas = _exp_avg_sq;
     _exp_avg = Tensor(_parameter.dtype(), _parameter.shape(), _parameter.device());
     _exp_avg_sq = Tensor(_parameter.dtype(), _parameter.shape(), _parameter.device());
diff --git a/madspace/src/driver/context.cpp b/madspace/src/driver/context.cpp
index 869c345fd0..ff20676882 100644
--- a/madspace/src/driver/context.cpp
+++ b/madspace/src/driver/context.cpp
@@ -370,7 +370,7 @@ Tensor Context::cached_tensor(std::size_t size) {
         }
     }
     std::size_t word_count = (size + 7) / 8;
-    Tensor new_tensor(DataType::dt_float, {word_count}, device());
+    Tensor new_tensor(DataType::dt_float64, {word_count}, device());
     tensors.insert(tensors.begin() + size_pos, new_tensor);
     return new_tensor;
 }
diff --git a/madspace/src/driver/discrete_optimizer.cpp b/madspace/src/driver/discrete_optimizer.cpp
index 2e12fa3a30..14f2116400 100644
--- a/madspace/src/driver/discrete_optimizer.cpp
+++ b/madspace/src/driver/discrete_optimizer.cpp
@@ -35,7 +35,7 @@ void DiscreteOptimizer::optimize() {
     //_sample_count = next_sample_count;
     for (auto [prob_name, data_item] : zip(_prob_names, _data)) {
         auto& [weight_sums, counts] = data_item;
-        Tensor prob(DataType::dt_float, _contexts.at(0)->global(prob_name).shape());
+        Tensor prob(DataType::dt_float64, _contexts.at(0)->global(prob_name).shape());
         auto prob_view = prob.view<double, 2>()[0];
 
         double norm = 0.;
diff --git a/madspace/src/driver/event_histograms.cpp b/madspace/src/driver/event_histograms.cpp
index a00376ee51..6af627b235 100644
--- a/madspace/src/driver/event_histograms.cpp
+++ b/madspace/src/driver/event_histograms.cpp
@@ -27,8 +27,9 @@ EventHistograms::EventHistograms(
 ) :
     _specs(specs),
     _reference_weight(
-        std::isfinite(reference_weight) && reference_weight != 0. ? reference_weight
-                                                                  : 1.
+        std::isfinite(reference_weight) && reference_weight != 0.
+            ? reference_weight
+            : 1.
     ) {
     if (!context) {
         throw std::invalid_argument("EventHistograms requires a context");
@@ -59,15 +60,19 @@ EventHistograms::EventHistograms(
                 "EventHistograms: one observable per histogram expected"
             );
         }
-        _runtimes.push_back(RuntimeData{
-            .runtime = build_runtime(obs->values.function(), context, false),
-            .particle_count = obs->particle_count,
-        });
+        _runtimes.push_back(
+            RuntimeData{
+                .runtime = build_runtime(obs->values.function(), context, false),
+                .particle_count = obs->particle_count,
+            }
+        );
     }
 }
 
 void EventHistograms::fill(
-    EventBuffer& buffer, const std::vector<double>& syst_weights, std::size_t weight_count
+    EventBuffer& buffer,
+    const std::vector<double>& syst_weights,
+    std::size_t weight_count
 ) {
     std::size_t count = buffer.event_count();
     if (count == 0) {
@@ -93,7 +98,7 @@ void EventHistograms::fill(
     for (auto& [subproc, indices] : groups) {
         auto& runtime_data = _runtimes.at(subproc).value();
         std::size_t pc = runtime_data.particle_count;
-        Tensor momenta(DataType::dt_float, {indices.size(), pc, 4});
+        Tensor momenta(DataType::dt_float64, {indices.size(), pc, 4});
         auto view = momenta.view<double, 3>();
         for (std::size_t n = 0; n < indices.size(); ++n) {
             for (std::size_t j = 0; j < pc; ++j) {
@@ -111,7 +116,7 @@ void EventHistograms::fill(
         }
         for (std::size_t o = 0; o < _specs.size(); ++o) {
             if (!_observable_index.at(o)) {
-                continue;  // filled from the event weight below
+                continue; // filled from the event weight below
             }
             Tensor out = outputs.at(_observable_index.at(o).value()).cpu().contiguous();
             if (out.shape().size() == 1) {
@@ -192,21 +197,22 @@ json EventHistograms::to_json(const SystematicsCalculator* systematics) const {
         return result;
     }
     double n = static_cast<double>(_event_count);
-    auto normalised = [&](const std::vector<double>& sums,
-                          const std::vector<double>& squares,
-                          std::vector<double>& values,
-                          std::vector<double>& errors) {
-        values.resize(sums.size());
-        errors.resize(sums.size());
-        for (std::size_t b = 0; b < sums.size(); ++b) {
-            // the event weights are "average" normalised: the cross section in a
-            // bin is the mean of (weight x indicator), its error the standard
-            // error of that mean
-            values[b] = sums[b] / n;
-            double var = (squares[b] - sums[b] * sums[b] / n) / n;
-            errors[b] = var > 0. ? std::sqrt(var / n) : 0.;
-        }
-    };
+    auto normalised =
+        [&](const std::vector<double>& sums,
+            const std::vector<double>& squares,
+            std::vector<double>& values,
+            std::vector<double>& errors) {
+            values.resize(sums.size());
+            errors.resize(sums.size());
+            for (std::size_t b = 0; b < sums.size(); ++b) {
+                // the event weights are "average" normalised: the cross section in a
+                // bin is the mean of (weight x indicator), its error the standard
+                // error of that mean
+                values[b] = sums[b] / n;
+                double var = (squares[b] - sums[b] * sums[b] / n) / n;
+                errors[b] = var > 0. ? std::sqrt(var / n) : 0.;
+            }
+        };
     std::vector<int> ids;
     std::vector<std::size_t> scale_indices;
     std::vector<PdfGroupInfo> pdf_groups;
@@ -217,7 +223,8 @@ json EventHistograms::to_json(const SystematicsCalculator* systematics) const {
     }
     for (std::size_t o = 0; o < _specs.size(); ++o) {
         auto& spec = _specs[o];
-        std::vector<std::vector<double>> values(_weight_count + 1), errors(_weight_count + 1);
+        std::vector<std::vector<double>> values(_weight_count + 1),
+            errors(_weight_count + 1);
         for (std::size_t c = 0; c <= _weight_count; ++c) {
             normalised(_sums[o][c], _square_sums[o][c], values[c], errors[c]);
         }
diff --git a/madspace/src/driver/io.cpp b/madspace/src/driver/io.cpp
index 87c5a79e5a..28c46c8aad 100644
--- a/madspace/src/driver/io.cpp
+++ b/madspace/src/driver/io.cpp
@@ -15,10 +15,13 @@ namespace {
 
 std::string dtype_to_str(DataType dtype) {
     switch (dtype) {
-    case DataType::dt_float:
+    case DataType::dt_float32:
+        return "<f4";
+        break;
+    case DataType::dt_float64:
         return "<f8";
         break;
-    case DataType::dt_int:
+    case DataType::dt_int32:
         return "<i4";
         break;
     default:
@@ -27,10 +30,12 @@ std::string dtype_to_str(DataType dtype) {
 }
 
 DataType str_to_dtype(std::string dtype) {
-    if (dtype == "<f8") {
-        return DataType::dt_float;
+    if (dtype == "<f4") {
+        return DataType::dt_float32;
+    } else if (dtype == "<f8") {
+        return DataType::dt_float64;
     } else if (dtype == "<i4") {
-        return DataType::dt_int;
+        return DataType::dt_int32;
     } else {
         throw std::invalid_argument("Unsupported data type");
     }
diff --git a/madspace/src/driver/systematics.cpp b/madspace/src/driver/systematics.cpp
index 9fd22e05ef..6a22c023f8 100644
--- a/madspace/src/driver/systematics.cpp
+++ b/madspace/src/driver/systematics.cpp
@@ -115,7 +115,9 @@ double SystematicsCalculator::dynamical_scale(
         return std::sqrt(std::max(0., e_tot * e_tot - pz_tot * pz_tot));
     }
     default:
-        throw std::invalid_argument(std::format("invalid dynamical scale choice {}", dyn));
+        throw std::invalid_argument(
+            std::format("invalid dynamical scale choice {}", dyn)
+        );
     }
 }
 
@@ -204,7 +206,9 @@ SystematicsCalculator::SystematicsCalculator(
         std::string prefix = std::format("{}.alpha_s{}", _prefix, index);
         grid.initialize_globals(_context, prefix);
         RunningCoupling coupling(grid, prefix);
-        _alpha_s_runtimes.push_back(build_runtime(coupling.function(), _context, false));
+        _alpha_s_runtimes.push_back(
+            build_runtime(coupling.function(), _context, false)
+        );
         return index;
     };
     add_alpha_s(nominal_alpha_s.value());
@@ -247,13 +251,15 @@ SystematicsCalculator::SystematicsCalculator(
                 remap = me_flavor_remap.at(i);
             }
             if (remap.size() < _subproc_args.at(i).beam_pdgs.size()) {
-                throw std::invalid_argument(std::format(
-                    "SystematicsCalculator: flavor remap of subprocess {} has {} "
-                    "entries, {} flavors expected",
-                    i,
-                    remap.size(),
-                    _subproc_args.at(i).beam_pdgs.size()
-                ));
+                throw std::invalid_argument(
+                    std::format(
+                        "SystematicsCalculator: flavor remap of subprocess {} has {} "
+                        "entries, {} flavors expected",
+                        i,
+                        remap.size(),
+                        _subproc_args.at(i).beam_pdgs.size()
+                    )
+                );
             }
             _matrix_elements.at(i) = MatrixElementData{
                 .runtime = build_runtime(me->function(), _context, false),
@@ -274,32 +280,39 @@ SystematicsCalculator::SystematicsCalculator(
     std::map<std::string, std::size_t> alpha_s_indices;
     for (auto& spec : _config.pdf_members) {
         if (!_config.has_pdf) {
-            _warnings.push_back(std::format(
-                "PDF variation {} member {} ignored: the beams have no PDF",
-                spec.set_name,
-                spec.member
-            ));
+            _warnings.push_back(
+                std::format(
+                    "PDF variation {} member {} ignored: the beams have no PDF",
+                    spec.set_name,
+                    spec.member
+                )
+            );
             continue;
         }
         if (_nominal_pdf2) {
             // a member of one set cannot stand in for two different sets
-            _warnings.push_back(std::format(
-                "PDF variation {} member {} ignored: the two beams use different "
-                "PDF sets",
-                spec.set_name,
-                spec.member
-            ));
+            _warnings.push_back(
+                std::format(
+                    "PDF variation {} member {} ignored: the two beams use different "
+                    "PDF sets",
+                    spec.set_name,
+                    spec.member
+                )
+            );
             continue;
         }
         bool same_set = spec.set_lhaid == _config.nominal_lhaid;
         if (!same_set && !_mur_supported) {
-            _warnings.push_back(std::format(
-                "PDF variation {} member {} ignored: the events have no uniform "
-                "alpha_s power and no matrix element to re-evaluate, so the alpha_s "
-                "of a different PDF set cannot be applied",
-                spec.set_name,
-                spec.member
-            ));
+            _warnings.push_back(
+                std::format(
+                    "PDF variation {} member {} ignored: the events have no uniform "
+                    "alpha_s power and no matrix element to re-evaluate, so the "
+                    "alpha_s "
+                    "of a different PDF set cannot be applied",
+                    spec.set_name,
+                    spec.member
+                )
+            );
             continue;
         }
         std::size_t alpha_s_index;
@@ -382,11 +395,10 @@ void SystematicsCalculator::build_variations() {
 
     // PDF variations: when members of the nominal set are requested, the group
     // starts with the nominal member itself so that the group is complete
-    bool has_nominal_set_members = std::any_of(
-        _members.begin(),
-        _members.end(),
-        [&](auto& spec) { return spec.set_lhaid == _config.nominal_lhaid; }
-    );
+    bool has_nominal_set_members =
+        std::any_of(_members.begin(), _members.end(), [&](auto& spec) {
+            return spec.set_lhaid == _config.nominal_lhaid;
+        });
     if (has_nominal_set_members) {
         _variations.push_back({id++, 1., 1., -1, -1});
     }
@@ -513,9 +525,9 @@ std::vector<double> SystematicsCalculator::evaluate_pdf(
     if (count == 0) {
         return values;
     }
-    Tensor x_tensor(DataType::dt_float, {count});
-    Tensor q_tensor(DataType::dt_float, {count});
-    Tensor slot_tensor(DataType::dt_int, {count});
+    Tensor x_tensor(DataType::dt_float64, {count});
+    Tensor q_tensor(DataType::dt_float64, {count});
+    Tensor slot_tensor(DataType::dt_int32, {count});
     auto x_view = x_tensor.view<double, 1>();
     auto q_view = q_tensor.view<double, 1>();
     auto slot_view = slot_tensor.view<me_int_t, 1>();
@@ -545,7 +557,7 @@ std::vector<double> SystematicsCalculator::evaluate_alpha_s(
     if (count == 0) {
         return values;
     }
-    Tensor q_tensor(DataType::dt_float, {count});
+    Tensor q_tensor(DataType::dt_float64, {count});
     auto q_view = q_tensor.view<double, 1>();
     for (std::size_t i = 0; i < count; ++i) {
         q_view[i] = q[i];
@@ -563,8 +575,9 @@ std::vector<double> SystematicsCalculator::evaluate_alpha_s(
     return values;
 }
 
-std::vector<std::array<double, 4>>
-SystematicsCalculator::event_momenta(EventBuffer& buffer, std::size_t event_index) const {
+std::vector<std::array<double, 4>> SystematicsCalculator::event_momenta(
+    EventBuffer& buffer, std::size_t event_index
+) const {
     std::vector<std::array<double, 4>> momenta;
     for (std::size_t j = 0; j < buffer.particle_count(); ++j) {
         auto particle = buffer.particle(event_index, j);
@@ -586,9 +599,9 @@ std::vector<double> SystematicsCalculator::matrix_elements(
     auto& me_data = _matrix_elements.at(subproc).value();
     std::size_t count = indices.size();
     std::size_t pc = me_data.particle_count;
-    Tensor momenta(DataType::dt_float, {count, pc, 4});
-    Tensor alpha_s_tensor(DataType::dt_float, {count});
-    Tensor flavor(DataType::dt_int, {count});
+    Tensor momenta(DataType::dt_float64, {count, pc, 4});
+    Tensor alpha_s_tensor(DataType::dt_float64, {count});
+    Tensor flavor(DataType::dt_int32, {count});
     auto mom_view = momenta.view<double, 3>();
     auto alpha_view = alpha_s_tensor.view<double, 1>();
     auto flavor_view = flavor.view<me_int_t, 1>();
@@ -621,8 +634,9 @@ std::vector<double> SystematicsCalculator::matrix_elements(
     return values;
 }
 
-LOReweightInfo
-SystematicsCalculator::reweight_info(EventBuffer& buffer, std::size_t event_index) const {
+LOReweightInfo SystematicsCalculator::reweight_info(
+    EventBuffer& buffer, std::size_t event_index
+) const {
     auto& layout = buffer.layout();
     bool has_beam1 = layout_has(layout, "fact_scale1");
     bool has_beam2 = layout_has(layout, "fact_scale2");
@@ -725,35 +739,37 @@ void SystematicsCalculator::compute(
     // x f(x, q) at the given points, point n evaluated with `second` if
     // on_second[n] and with `first` otherwise. One batched call when both are
     // the same grid.
-    auto evaluate_beams = [&](const PdfEvaluator& first,
-                              const PdfEvaluator& second,
-                              const std::vector<double>& x,
-                              const std::vector<double>& q,
-                              const std::vector<me_int_t>& slots,
-                              const std::vector<char>& on_second) {
-        if (&first == &second) {
-            return evaluate_pdf(first, x, q, slots);
-        }
-        std::array<std::vector<double>, 2> grid_x, grid_q;
-        std::array<std::vector<me_int_t>, 2> grid_slots;
-        for (std::size_t n = 0; n < x.size(); ++n) {
-            grid_x[on_second[n]].push_back(x[n]);
-            grid_q[on_second[n]].push_back(q[n]);
-            grid_slots[on_second[n]].push_back(slots[n]);
-        }
-        std::array<std::vector<double>, 2> grid_values{
-            evaluate_pdf(first, grid_x[0], grid_q[0], grid_slots[0]),
-            evaluate_pdf(second, grid_x[1], grid_q[1], grid_slots[1]),
+    auto evaluate_beams =
+        [&](const PdfEvaluator& first,
+            const PdfEvaluator& second,
+            const std::vector<double>& x,
+            const std::vector<double>& q,
+            const std::vector<me_int_t>& slots,
+            const std::vector<char>& on_second) {
+            if (&first == &second) {
+                return evaluate_pdf(first, x, q, slots);
+            }
+            std::array<std::vector<double>, 2> grid_x, grid_q;
+            std::array<std::vector<me_int_t>, 2> grid_slots;
+            for (std::size_t n = 0; n < x.size(); ++n) {
+                grid_x[on_second[n]].push_back(x[n]);
+                grid_q[on_second[n]].push_back(q[n]);
+                grid_slots[on_second[n]].push_back(slots[n]);
+            }
+            std::array<std::vector<double>, 2> grid_values{
+                evaluate_pdf(first, grid_x[0], grid_q[0], grid_slots[0]),
+                evaluate_pdf(second, grid_x[1], grid_q[1], grid_slots[1]),
+            };
+            std::vector<double> values(x.size());
+            std::array<std::size_t, 2> next{0, 0};
+            for (std::size_t n = 0; n < x.size(); ++n) {
+                values[n] = grid_values[on_second[n]][next[on_second[n]]++];
+            }
+            return values;
         };
-        std::vector<double> values(x.size());
-        std::array<std::size_t, 2> next{0, 0};
-        for (std::size_t n = 0; n < x.size(); ++n) {
-            values[n] = grid_values[on_second[n]][next[on_second[n]]++];
-        }
-        return values;
-    };
-    const PdfEvaluator* nominal_pdf2 =
-        _nominal_pdf2 ? &_nominal_pdf2.value() : (_nominal_pdf ? &_nominal_pdf.value() : nullptr);
+    const PdfEvaluator* nominal_pdf2 = _nominal_pdf2
+        ? &_nominal_pdf2.value()
+        : (_nominal_pdf ? &_nominal_pdf.value() : nullptr);
 
     // Nominal PDF product when the events do not carry it: one batched call
     if (use_pdf && !has_partial) {
@@ -770,9 +786,8 @@ void SystematicsCalculator::compute(
                 on_second.push_back(!in.swapped);
             }
         }
-        auto values = evaluate_beams(
-            _nominal_pdf.value(), *nominal_pdf2, x, q, slots, on_second
-        );
+        auto values =
+            evaluate_beams(_nominal_pdf.value(), *nominal_pdf2, x, q, slots, on_second);
         for (std::size_t i = 0, n = 0; i < count; ++i) {
             double product = 1.;
             if (has_beam1) {
@@ -789,46 +804,47 @@ void SystematicsCalculator::compute(
     // (nominal grid: every scale variation; member grids: the nominal scales)
     std::vector<std::vector<double>> r_pdf(count, std::vector<double>(var_count, 1.));
     if (use_pdf) {
-        auto evaluate_variations = [&](const PdfEvaluator& evaluator,
-                                       const PdfEvaluator& evaluator2,
-                                       const std::vector<std::size_t>& var_indices) {
-            std::vector<double> x, q;
-            std::vector<me_int_t> slots;
-            std::vector<char> on_second;
-            for (auto& in : inputs) {
-                for (std::size_t k : var_indices) {
-                    auto& var = _variations.at(k);
-                    if (has_beam1) {
-                        x.push_back(in.x1);
-                        q.push_back(var.muf * scale_of(in, var.dyn, in.muf1));
-                        slots.push_back(in.slot1);
-                        on_second.push_back(in.swapped);
-                    }
-                    if (has_beam2) {
-                        x.push_back(in.x2);
-                        q.push_back(var.muf * scale_of(in, var.dyn, in.muf2));
-                        slots.push_back(in.slot2);
-                        on_second.push_back(!in.swapped);
+        auto evaluate_variations =
+            [&](const PdfEvaluator& evaluator,
+                const PdfEvaluator& evaluator2,
+                const std::vector<std::size_t>& var_indices) {
+                std::vector<double> x, q;
+                std::vector<me_int_t> slots;
+                std::vector<char> on_second;
+                for (auto& in : inputs) {
+                    for (std::size_t k : var_indices) {
+                        auto& var = _variations.at(k);
+                        if (has_beam1) {
+                            x.push_back(in.x1);
+                            q.push_back(var.muf * scale_of(in, var.dyn, in.muf1));
+                            slots.push_back(in.slot1);
+                            on_second.push_back(in.swapped);
+                        }
+                        if (has_beam2) {
+                            x.push_back(in.x2);
+                            q.push_back(var.muf * scale_of(in, var.dyn, in.muf2));
+                            slots.push_back(in.slot2);
+                            on_second.push_back(!in.swapped);
+                        }
                     }
                 }
-            }
-            auto values =
-                evaluate_beams(evaluator, evaluator2, x, q, slots, on_second);
-            for (std::size_t i = 0, n = 0; i < count; ++i) {
-                for (std::size_t k : var_indices) {
-                    double product = 1.;
-                    if (has_beam1) {
-                        product *= values[n++];
-                    }
-                    if (has_beam2) {
-                        product *= values[n++];
+                auto values =
+                    evaluate_beams(evaluator, evaluator2, x, q, slots, on_second);
+                for (std::size_t i = 0, n = 0; i < count; ++i) {
+                    for (std::size_t k : var_indices) {
+                        double product = 1.;
+                        if (has_beam1) {
+                            product *= values[n++];
+                        }
+                        if (has_beam2) {
+                            product *= values[n++];
+                        }
+                        r_pdf[i][k] = inputs[i].nominal_product == 0.
+                            ? 0.
+                            : product / inputs[i].nominal_product;
                     }
-                    r_pdf[i][k] = inputs[i].nominal_product == 0.
-                        ? 0.
-                        : product / inputs[i].nominal_product;
                 }
-            }
-        };
+            };
         std::vector<std::size_t> nominal_grid_vars;
         std::vector<std::vector<std::size_t>> member_vars(_member_pdfs.size());
         for (std::size_t k = 0; k < var_count; ++k) {
@@ -909,7 +925,8 @@ void SystematicsCalculator::compute(
             if (in.use_me || in.qcd_power == 0) {
                 continue;
             }
-            ratio[i] = in.alpha_s == 0. ? 0. : std::pow(varied[i] / in.alpha_s, in.qcd_power);
+            ratio[i] =
+                in.alpha_s == 0. ? 0. : std::pow(varied[i] / in.alpha_s, in.qcd_power);
         }
         for (auto& [subproc, indices] : me_events) {
             std::vector<double> alpha_s(indices.size());
@@ -1067,8 +1084,7 @@ json SystematicsCalculator::summary() const {
     json result{
         {"variations", variations},
         {"nominal",
-         {{"pdf_set", _config.nominal_set_name},
-          {"pdf_lhaid", _config.nominal_lhaid}}},
+         {{"pdf_set", _config.nominal_set_name}, {"pdf_lhaid", _config.nominal_lhaid}}},
         {"event_count", _event_count},
         {"warnings", _warnings},
     };
diff --git a/madspace/src/driver/tensor.cpp b/madspace/src/driver/tensor.cpp
index 6842b1686d..0c38156485 100644
--- a/madspace/src/driver/tensor.cpp
+++ b/madspace/src/driver/tensor.cpp
@@ -1,5 +1,6 @@
 #include "madspace/driver/tensor.hpp"
 
+#include <algorithm>
 #include <variant>
 
 using namespace madspace;
@@ -231,6 +232,44 @@ std::size_t Tensor::init_stride() {
     return stride_prod * dtype_size();
 }
 
+Tensor Tensor::cast(DataType dtype) const {
+    Tensor source = cpu().contiguous();
+    if (source.dtype() == dtype) {
+        return source;
+    }
+    Tensor target(dtype, source.shape());
+    std::size_t count = source.shape().product();
+    auto convert = [&]<typename From>(From* from) {
+        switch (dtype) {
+        case DataType::dt_int32:
+            std::copy_n(from, count, static_cast<me_int_t*>(target.data()));
+            break;
+        case DataType::dt_float32:
+            std::copy_n(from, count, static_cast<float*>(target.data()));
+            break;
+        case DataType::dt_float64:
+            std::copy_n(from, count, static_cast<double*>(target.data()));
+            break;
+        default:
+            throw std::invalid_argument("invalid dtype in cast");
+        }
+    };
+    switch (source.dtype()) {
+    case DataType::dt_int32:
+        convert(static_cast<me_int_t*>(source.data()));
+        break;
+    case DataType::dt_float32:
+        convert(static_cast<float*>(source.data()));
+        break;
+    case DataType::dt_float64:
+        convert(static_cast<double*>(source.data()));
+        break;
+    default:
+        throw std::invalid_argument("invalid dtype in cast");
+    }
+    return target;
+}
+
 Tensor Tensor::simd_padded_literal(
     const LiteralValue& value, std::size_t lanes, DevicePtr device
 ) {
@@ -275,8 +314,7 @@ Tensor Tensor::simd_padded_literal(
     // element j of lane l is stored at l + lanes * j
     auto [dtype, buffer] = std::visit(
         [&]<typename T>(const std::vector<T>& vec) {
-            auto dtype =
-                std::is_same_v<T, me_int_t> ? DataType::dt_int : DataType::dt_float;
+            auto dtype = dtype_of<T>();
             std::vector<T> padded(count * lanes);
             for (std::size_t j = 0; j < count; ++j) {
                 std::fill_n(padded.begin() + j * lanes, lanes, vec[j]);
@@ -294,7 +332,16 @@ Tensor Tensor::simd_padded_literal(
         stride_prod *= shape[i];
     }
     Tensor ret(new TensorImpl{
-        dtype, shape, device, buffer.data(), false, std::nullopt, buffer.impl, 1, stride, 1
+        dtype,
+        shape,
+        device,
+        buffer.data(),
+        false,
+        std::nullopt,
+        buffer.impl,
+        1,
+        stride,
+        1
     });
     ret.impl->simd_padded = true;
     // non-negative integer scalars double as batch sizes, see the scalar constructor
diff --git a/madspace/src/gpu/device.cu b/madspace/src/gpu/device.cu
index 745e22f611..3b19d64a56 100644
--- a/madspace/src/gpu/device.cu
+++ b/madspace/src/gpu/device.cu
@@ -13,6 +13,50 @@ void wait_for_tensor(const Tensor& tensor) {
         check_error(gpuStreamSynchronize(reinterpret_cast<gpuStream_t>(*stream)));
     }
 }
+
+// calls func with the kernel policy of a float dtype
+template <typename Func>
+void dispatch_float(DataType dtype, const char* op, Func&& func) {
+    switch (dtype) {
+    case DataType::dt_float64:
+        func.template operator()<GpuTypes>();
+        break;
+    case DataType::dt_float32:
+        func.template operator()<GpuTypes32>();
+        break;
+    default:
+        throw std::runtime_error(std::string("invalid dtype in ") + op);
+    }
+}
+
+// sums a contiguous (batch_size, n) source over the batch dimension
+template <typename F>
+__global__ void
+kernel_batch_sum(std::size_t batch_size, std::size_t n, const F* source, F* output) {
+    std::size_t i = blockDim.x * blockIdx.x + threadIdx.x;
+    if (i >= n) {
+        return;
+    }
+    F sum = 0;
+    for (std::size_t j = 0; j < batch_size; ++j) {
+        sum += source[j * n + i];
+    }
+    output[i] = sum;
+}
+
+template <typename F>
+void launch_batch_sum(const Tensor& source, Tensor& sum, gpuStream_t stream) {
+    std::size_t n = sum.shape().product();
+    launch_kernel(
+        kernel_batch_sum<F>,
+        n,
+        stream,
+        source.size(0),
+        n,
+        static_cast<const F*>(source.data()),
+        static_cast<F*>(sum.data())
+    );
+}
 } // namespace
 
 std::pair<void*, Tensor> GpuDevice::allocate(std::size_t size, AllocHint hint) const {
@@ -107,18 +151,21 @@ void GpuDevice::adam_step(
     wait_for_tensor(exp_avg);
     wait_for_tensor(exp_avg_sq);
     AsyncGpuDevice device(*this, gpuStreamPerThread, 0);
-    tensor_foreach_dynamic<kernel_adam_step<GpuTypes>, 1, 3>(
-        {&gradient},
-        {&parameter, &exp_avg, &exp_avg_sq},
-        gradient.size(0),
-        device,
-        step_size,
-        beta1,
-        beta2,
-        eps,
-        bias_corr2_sqrt,
-        weight_decay
-    );
+    dispatch_float(parameter.dtype(), "adam_step", [&]<typename T>() {
+        using F = typename T::FVal;
+        tensor_foreach_dynamic<kernel_adam_step<T>, 1, 3>(
+            {&gradient},
+            {&parameter, &exp_avg, &exp_avg_sq},
+            gradient.size(0),
+            device,
+            static_cast<F>(step_size),
+            static_cast<F>(beta1),
+            static_cast<F>(beta2),
+            static_cast<F>(eps),
+            static_cast<F>(bias_corr2_sqrt),
+            static_cast<F>(weight_decay)
+        );
+    });
     check_error(gpuStreamSynchronize(gpuStreamPerThread));
 }
 
@@ -147,7 +194,8 @@ MemPool::MemPool(
             pool.parent_tensor = parent_tensor;
             pool.capacity = parent_tensor.byte_size();
         } else {
-            pool.parent_tensor = Tensor(DataType::dt_float, {word_count}, async_device);
+            pool.parent_tensor =
+                Tensor(DataType::dt_float64, {word_count}, async_device);
             pool.capacity = word_count * 8;
         }
         pool.needed_size = word_count * 8;
@@ -307,46 +355,65 @@ void AsyncGpuDevice::memcpy(void* to, void* from, std::size_t size) const {
 }
 
 void AsyncGpuDevice::tensor_copy(const Tensor& source, Tensor& target) const {
-    if (source.dtype() == DataType::dt_float && target.dtype() == DataType::dt_float) {
-        tensor_foreach_dynamic<kernel_copy<GpuTypes>, 1, 1>(
-            {&source}, {&target}, target.size(0), *this
-        );
-    } else if (source.dtype() == DataType::dt_int &&
-               target.dtype() == DataType::dt_int) {
+    if (source.dtype() != target.dtype()) {
+        throw std::runtime_error("invalid dtype in copy");
+    }
+    if (source.dtype() == DataType::dt_int32) {
         tensor_foreach_dynamic<kernel_copy_int<GpuTypes>, 1, 1>(
             {&source}, {&target}, target.size(0), *this
         );
-    } else {
-        throw std::runtime_error("invalid dtype in copy");
+        return;
     }
+    dispatch_float(source.dtype(), "copy", [&]<typename T>() {
+        tensor_foreach_dynamic<kernel_copy<T>, 1, 1>(
+            {&source}, {&target}, target.size(0), *this
+        );
+    });
 }
 
 void AsyncGpuDevice::tensor_zero(Tensor& tensor) const {
-    if (tensor.dtype() == DataType::dt_float) {
-        if (tensor.is_contiguous()) {
-            check_error(gpuMemsetAsync(tensor.data(), 0, tensor.byte_size(), _stream));
-        } else {
-            tensor_foreach_dynamic<kernel_zero<GpuTypes>, 1, 1>(
-                {&tensor}, {&tensor}, tensor.size(0), *this
-            );
-        }
-    } else if (tensor.dtype() == DataType::dt_int) {
-        if (tensor.is_contiguous()) {
-            check_error(gpuMemsetAsync(tensor.data(), 0, tensor.byte_size(), _stream));
-        } else {
-            tensor_foreach_dynamic<kernel_zero_int<GpuTypes>, 1, 1>(
-                {&tensor}, {&tensor}, tensor.size(0), *this
-            );
-        }
-    } else {
+    if (tensor.dtype() == DataType::batch_sizes) {
         throw std::runtime_error("invalid dtype in zero");
     }
+    if (tensor.is_contiguous()) {
+        check_error(gpuMemsetAsync(tensor.data(), 0, tensor.byte_size(), _stream));
+        return;
+    }
+    if (tensor.dtype() == DataType::dt_int32) {
+        tensor_foreach_dynamic<kernel_zero_int<GpuTypes>, 1, 1>(
+            {&tensor}, {&tensor}, tensor.size(0), *this
+        );
+        return;
+    }
+    dispatch_float(tensor.dtype(), "zero", [&]<typename T>() {
+        tensor_foreach_dynamic<kernel_zero<T>, 1, 1>(
+            {&tensor}, {&tensor}, tensor.size(0), *this
+        );
+    });
 }
 
 void AsyncGpuDevice::tensor_add(const Tensor& source, Tensor& target) const {
-    tensor_foreach_dynamic<kernel_add_inplace<GpuTypes>, 1, 1>(
-        {&source}, {&target}, target.size(0), *this
-    );
+    if (source.dtype() != target.dtype()) {
+        throw std::runtime_error("invalid dtype in add");
+    }
+    if (target.size(0) == 1 && source.size(0) > 1) {
+        // a broadcast target (e.g. the gradient of a global) accumulates the
+        // whole source batch, as on the CPU
+        auto source_contig = source.contiguous(*this, AllocHint::temporary);
+        Tensor sum(target.dtype(), target.shape(), *this, AllocHint::temporary);
+        dispatch_float(source.dtype(), "add", [&]<typename T>() {
+            launch_batch_sum<typename T::FVal>(source_contig, sum, _stream);
+        });
+        tensor_add(sum, target);
+        sum.reset(*this);
+        source_contig.reset(*this);
+        return;
+    }
+    dispatch_float(source.dtype(), "add", [&]<typename T>() {
+        tensor_foreach_dynamic<kernel_add_inplace<T>, 1, 1>(
+            {&source}, {&target}, target.size(0), *this
+        );
+    });
 }
 
 void AsyncGpuDevice::tensor_cpu(const Tensor& source, Tensor& target) const {
diff --git a/madspace/src/gpu/gpu_abstraction.cuh b/madspace/src/gpu/gpu_abstraction.cuh
index 5c62d5f700..5d71f2c89e 100644
--- a/madspace/src/gpu/gpu_abstraction.cuh
+++ b/madspace/src/gpu/gpu_abstraction.cuh
@@ -45,6 +45,8 @@
 #define gpublasSetStream cublasSetStream
 #define gpublasDgemm cublasDgemm
 #define gpublasDgemv cublasDgemv
+#define gpublasSgemm cublasSgemm
+#define gpublasSgemv cublasSgemv
 #define GPUBLAS_OP_N CUBLAS_OP_N
 #define GPUBLAS_OP_T CUBLAS_OP_T
 
@@ -95,6 +97,8 @@
 #define gpublasSetStream rocblas_set_stream
 #define gpublasDgemm rocblas_dgemm
 #define gpublasDgemv rocblas_dgemv
+#define gpublasSgemm rocblas_sgemm
+#define gpublasSgemv rocblas_sgemv
 #define GPUBLAS_OP_N rocblas_operation_none
 #define GPUBLAS_OP_T rocblas_operation_transpose
 
diff --git a/madspace/src/gpu/kernel_definitions.cuh b/madspace/src/gpu/kernel_definitions.cuh
index a4f8c1c03b..721c02bddd 100644
--- a/madspace/src/gpu/kernel_definitions.cuh
+++ b/madspace/src/gpu/kernel_definitions.cuh
@@ -8,28 +8,37 @@
 namespace madspace {
 namespace kernels {
 
-struct GpuTypes {
+template <typename F>
+struct GpuTypesT {
     template <int dim>
-    using FIn = const gpu::GpuTensorView<double, dim>;
+    using FIn = const gpu::GpuTensorView<F, dim>;
     template <int dim>
     using IIn = const gpu::GpuTensorView<me_int_t, dim>;
     template <int dim>
-    using FOut = gpu::GpuTensorView<double, dim>;
+    using FOut = gpu::GpuTensorView<F, dim>;
     template <int dim>
     using IOut = gpu::GpuTensorView<me_int_t, dim>;
     // single inputs are broadcast through a zero batch stride on the GPU
     template <int dim>
-    using FSIn = const gpu::GpuTensorView<double, dim>;
+    using FSIn = const gpu::GpuTensorView<F, dim>;
     template <int dim>
     using ISIn = const gpu::GpuTensorView<me_int_t, dim>;
-    using FVal = double;
+    using FVal = F;
     using IVal = me_int_t;
     using BVal = bool;
+    // the same policy with a different float type, for conversion kernels
+    using F32 = GpuTypesT<float>;
+    using F64 = GpuTypesT<double>;
 };
+using GpuTypes = GpuTypesT<double>;
+using GpuTypes32 = GpuTypesT<float>;
 
 inline __device__ double where(bool condition, double val_true, double val_false) {
     return condition ? val_true : val_false;
 }
+inline __device__ float where(bool condition, float val_true, float val_false) {
+    return condition ? val_true : val_false;
+}
 inline __device__ me_int_t
 where(bool condition, me_int_t val_true, me_int_t val_false) {
     return condition ? val_true : val_false;
@@ -40,6 +49,12 @@ inline __device__ double min(double arg1, double arg2) {
 inline __device__ double max(double arg1, double arg2) {
     return arg1 > arg2 ? arg1 : arg2;
 }
+inline __device__ float min(float arg1, float arg2) {
+    return arg1 < arg2 ? arg1 : arg2;
+}
+inline __device__ float max(float arg1, float arg2) {
+    return arg1 > arg2 ? arg1 : arg2;
+}
 inline __device__ std::size_t single_index(me_int_t arg) { return arg; }
 
 using ::acos;
diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu
index 1b2624b72f..b2741e55d1 100644
--- a/madspace/src/gpu/runtime.cu
+++ b/madspace/src/gpu/runtime.cu
@@ -125,6 +125,25 @@ void op_matrix_element(
     }
 }
 
+template <typename F, typename... Args>
+auto gpublas_gemm(Args... args) {
+    if constexpr (std::is_same_v<F, float>) {
+        return gpublasSgemm(args...);
+    } else {
+        return gpublasDgemm(args...);
+    }
+}
+
+template <typename F, typename... Args>
+auto gpublas_gemv(Args... args) {
+    if constexpr (std::is_same_v<F, float>) {
+        return gpublasSgemv(args...);
+    } else {
+        return gpublasDgemv(args...);
+    }
+}
+
+template <typename F>
 void op_matmul(
     const GpuRuntime::Instruction& instruction,
     TensorVec& locals,
@@ -141,10 +160,7 @@ void op_matmul(
     std::size_t dims_in = input.size(1);
     std::size_t dims_out = weight.size(1);
     output = Tensor(
-        DataType::dt_float,
-        {batch_size, dims_out},
-        device,
-        instruction.output_alloc_hints[0]
+        dtype_of<F>(), {batch_size, dims_out}, device, instruction.output_alloc_hints[0]
     );
     output.copy_from(bias, device);
     if (batch_size == 0) {
@@ -153,36 +169,39 @@ void op_matmul(
 
     gpublasHandle_t handle = instruction.runtime.gpublas_handle();
     check_error(gpublasSetStream(handle, device.stream()));
-    double alpha = 1., beta = 1.;
-    check_error(gpublasDgemm(
-        handle,
-        GPUBLAS_OP_N,
-        GPUBLAS_OP_T,
-        batch_size,
-        dims_out,
-        dims_in,
-        &alpha,
-        static_cast<double*>(input.data()),
-        batch_size,
-        static_cast<double*>(weight.data()),
-        dims_out,
-        &beta,
-        static_cast<double*>(output.data()),
-        batch_size
-    ));
+    F alpha = 1., beta = 1.;
+    check_error(
+        gpublas_gemm<F>(
+            handle,
+            GPUBLAS_OP_N,
+            GPUBLAS_OP_T,
+            batch_size,
+            dims_out,
+            dims_in,
+            &alpha,
+            static_cast<F*>(input.data()),
+            batch_size,
+            static_cast<F*>(weight.data()),
+            dims_out,
+            &beta,
+            static_cast<F*>(output.data()),
+            batch_size
+        )
+    );
     input.reset(device);
     weight.reset(device);
     bias.reset(device);
 }
 
-__global__ void
-kernel_one(std::size_t batch_size, GpuTensorView<double, 1, true> output) {
+template <typename F>
+__global__ void kernel_one(std::size_t batch_size, GpuTensorView<F, 1, true> output) {
     me_int_t i = blockDim.x * blockIdx.x + threadIdx.x;
     if (i < batch_size) {
         output[i] = 1.;
     }
 }
 
+template <typename F>
 void backward_op_matmul(
     const GpuRuntime::Instruction& instruction,
     TensorVec& locals,
@@ -205,30 +224,18 @@ void backward_op_matmul(
 
     if (!input_grad) {
         input_grad = Tensor(
-            DataType::dt_float,
-            input.shape(),
-            device,
-            instruction.input_grad_alloc_hints[0]
+            dtype_of<F>(), input.shape(), device, instruction.input_grad_alloc_hints[0]
         );
-        // input_grad.zero(device);
     }
     if (!weight_grad) {
         weight_grad = Tensor(
-            DataType::dt_float,
-            weight.shape(),
-            device,
-            instruction.input_grad_alloc_hints[1]
+            dtype_of<F>(), weight.shape(), device, instruction.input_grad_alloc_hints[1]
         );
-        // weight_grad.zero(device);
     }
     if (!bias_grad) {
         bias_grad = Tensor(
-            DataType::dt_float,
-            {1, dims_out},
-            device,
-            instruction.input_grad_alloc_hints[2]
+            dtype_of<F>(), {1, dims_out}, device, instruction.input_grad_alloc_hints[2]
         );
-        // bias_grad.zero(device);
     }
     if (batch_size == 0) {
         input.reset(device);
@@ -237,66 +244,72 @@ void backward_op_matmul(
         return;
     }
 
-    double alpha = 1., beta = 1.;
+    F alpha = 1., beta = 1.;
     gpublasHandle_t handle = instruction.runtime.gpublas_handle();
     gpuStream_t stream = device.stream();
     check_error(gpublasSetStream(handle, stream));
 
     // compute input_grad += output_grad * weight
-    check_error(gpublasDgemm(
-        handle,
-        GPUBLAS_OP_N,
-        GPUBLAS_OP_N,
-        batch_size,
-        dims_in,
-        dims_out,
-        &alpha,
-        static_cast<double*>(output_grad.data()),
-        batch_size,
-        static_cast<double*>(weight.data()),
-        dims_out,
-        &beta,
-        static_cast<double*>(input_grad.data()),
-        batch_size
-    ));
+    check_error(
+        gpublas_gemm<F>(
+            handle,
+            GPUBLAS_OP_N,
+            GPUBLAS_OP_N,
+            batch_size,
+            dims_in,
+            dims_out,
+            &alpha,
+            static_cast<F*>(output_grad.data()),
+            batch_size,
+            static_cast<F*>(weight.data()),
+            dims_out,
+            &beta,
+            static_cast<F*>(input_grad.data()),
+            batch_size
+        )
+    );
 
     // compute weight_grad += output_grad.T * input
-    check_error(gpublasDgemm(
-        handle,
-        GPUBLAS_OP_T,
-        GPUBLAS_OP_N,
-        dims_out,
-        dims_in,
-        batch_size,
-        &alpha,
-        static_cast<double*>(output_grad.data()),
-        batch_size,
-        static_cast<double*>(input.data()),
-        batch_size,
-        &beta,
-        static_cast<double*>(weight_grad.data()),
-        dims_out
-    ));
+    check_error(
+        gpublas_gemm<F>(
+            handle,
+            GPUBLAS_OP_T,
+            GPUBLAS_OP_N,
+            dims_out,
+            dims_in,
+            batch_size,
+            &alpha,
+            static_cast<F*>(output_grad.data()),
+            batch_size,
+            static_cast<F*>(input.data()),
+            batch_size,
+            &beta,
+            static_cast<F*>(weight_grad.data()),
+            dims_out
+        )
+    );
 
     // compute bias_grad += sum_i output_grad_ij
-    Tensor ones(DataType::dt_float, {batch_size}, device, AllocHint::temporary);
+    Tensor ones(dtype_of<F>(), {batch_size}, device, AllocHint::temporary);
     launch_kernel(
-        kernel_one, batch_size, device.stream(), batch_size, ones.view<double, 1>()
+        kernel_one<F>, batch_size, device.stream(), batch_size, ones.view<F, 1>()
+    );
+    check_error(
+        gpublas_gemv<F>(
+            handle,
+            GPUBLAS_OP_T,
+            batch_size,
+            dims_out,
+            &alpha,
+            static_cast<F*>(output_grad.data()),
+            batch_size,
+            static_cast<F*>(ones.data()),
+            1,
+            &beta,
+            static_cast<F*>(bias_grad.data()),
+            1
+        )
     );
-    check_error(gpublasDgemv(
-        handle,
-        GPUBLAS_OP_T,
-        batch_size,
-        dims_out,
-        &alpha,
-        static_cast<double*>(output_grad.data()),
-        batch_size,
-        static_cast<double*>(ones.data()),
-        1,
-        &beta,
-        static_cast<double*>(bias_grad.data()),
-        1
-    ));
     ones.reset(device);
     input.reset(device);
     weight.reset(device);
@@ -326,9 +339,9 @@ void op_nonzero(
     auto& input = locals[instruction.input_indices[0]];
     auto batch_size = input.size(0);
     auto& output = locals[instruction.output_indices[0]];
-    Tensor indices_tmp(DataType::dt_int, {batch_size}, device, AllocHint::temporary);
+    Tensor indices_tmp(DataType::dt_int32, {batch_size}, device, AllocHint::temporary);
     Tensor output_tmp(
-        DataType::dt_int, {batch_size}, device, instruction.output_alloc_hints[0]
+        DataType::dt_int32, {batch_size}, device, instruction.output_alloc_hints[0]
     );
     launch_kernel(
         kernel_nonzero,
@@ -397,8 +410,8 @@ void batch_gather_impl(
     Sizes out_shape = values.shape();
     out_shape[0] = batch_size;
 
-    if (values.dtype() == DataType::dt_float) {
-        selection = Tensor(DataType::dt_float, out_shape, device, hint);
+    if (values.dtype() == DataType::dt_float64) {
+        selection = Tensor(DataType::dt_float64, out_shape, device, hint);
         launch_kernel(
             batch_gather_kernel<dim>,
             batch_size,
@@ -408,8 +421,8 @@ void batch_gather_impl(
             values.view<double, dim>(),
             selection.view<double, dim>()
         );
-    } else if (values.dtype() == DataType::dt_int) {
-        selection = Tensor(DataType::dt_int, out_shape, device, hint);
+    } else if (values.dtype() == DataType::dt_int32) {
+        selection = Tensor(DataType::dt_int32, out_shape, device, hint);
         launch_kernel(
             batch_gather_kernel_int<dim>,
             batch_size,
@@ -483,7 +496,7 @@ template <int dim>
 void batch_scatter_impl(
     Tensor& indices, Tensor& source, Tensor& output, const AsyncGpuDevice& device
 ) {
-    if (source.dtype() == DataType::dt_float) {
+    if (source.dtype() == DataType::dt_float64) {
         auto batch_size = indices.size(0);
         launch_kernel(
             batch_scatter_kernel<dim>,
@@ -494,7 +507,7 @@ void batch_scatter_impl(
             source.view<double, dim>(),
             output.view<double, dim>()
         );
-    } else if (source.dtype() == DataType::dt_int) {
+    } else if (source.dtype() == DataType::dt_int32) {
         auto batch_size = indices.size(0);
         launch_kernel(
             batch_scatter_kernel_int<dim>,
@@ -558,14 +571,14 @@ void op_batch_split_by_index(
     if (batch_size == 0) {
         for (std::size_t k = 0; k < count; ++k) {
             locals[instruction.output_indices[k]] = Tensor(
-                DataType::dt_int, {0}, device, instruction.output_alloc_hints[k]
+                DataType::dt_int32, {0}, device, instruction.output_alloc_hints[k]
             );
         }
         return;
     }
 
     // per-bucket counts, needed on the host to size and split the outputs below
-    Tensor sizes(DataType::dt_int, {count}, device, AllocHint::temporary);
+    Tensor sizes(DataType::dt_int32, {count}, device, AllocHint::temporary);
     {
         std::size_t temp_storage_bytes = 0;
         check_error(
@@ -582,7 +595,7 @@ void op_batch_split_by_index(
             )
         );
         Tensor temp(
-            DataType::dt_float,
+            DataType::dt_float64,
             {(temp_storage_bytes + 7) / 8},
             device,
             AllocHint::temporary
@@ -605,15 +618,15 @@ void op_batch_split_by_index(
 
     // stable sort of the original positions by bucket index: afterwards each
     // bucket's positions are contiguous, in their original relative order
-    Tensor values_in(DataType::dt_int, {batch_size}, device, AllocHint::temporary);
+    Tensor values_in(DataType::dt_int32, {batch_size}, device, AllocHint::temporary);
     auto values_in_ptr =
         thrust::device_pointer_cast(static_cast<me_int_t*>(values_in.data()));
     thrust::sequence(
         thrust_par.on(device.stream()), values_in_ptr, values_in_ptr + batch_size
     );
-    Tensor keys_out(DataType::dt_int, {batch_size}, device, AllocHint::temporary);
+    Tensor keys_out(DataType::dt_int32, {batch_size}, device, AllocHint::temporary);
     Tensor values_out(
-        DataType::dt_int, {batch_size}, device, instruction.output_alloc_hints[0]
+        DataType::dt_int32, {batch_size}, device, instruction.output_alloc_hints[0]
     );
     {
         std::size_t temp_storage_bytes = 0;
@@ -632,7 +645,7 @@ void op_batch_split_by_index(
             )
         );
         Tensor temp(
-            DataType::dt_float,
+            DataType::dt_float64,
             {(temp_storage_bytes + 7) / 8},
             device,
             AllocHint::temporary
@@ -718,7 +731,7 @@ void batch_reduce_sum_mean_impl(
     AllocHint hint = instruction.output_alloc_hints[0];
     std::size_t batch_size = input.size(0);
 
-    Tensor out(DataType::dt_float, {1}, device, hint);
+    Tensor out(DataType::dt_float64, {1}, device, hint);
     std::size_t temp_storage_bytes = 0;
     check_error(
         cub::DeviceReduce::Sum(
@@ -731,7 +744,10 @@ void batch_reduce_sum_mean_impl(
         )
     );
     Tensor temp(
-        DataType::dt_float, {(temp_storage_bytes + 7) / 8}, device, AllocHint::temporary
+        DataType::dt_float64,
+        {(temp_storage_bytes + 7) / 8},
+        device,
+        AllocHint::temporary
     );
     check_error(
         cub::DeviceReduce::Sum(
@@ -774,7 +790,7 @@ void batch_reduce_sum_mean_backward_impl(
     auto& input_grad = local_grads[instruction.input_indices[0]];
     if (!input_grad) {
         input_grad = Tensor(
-            DataType::dt_float,
+            DataType::dt_float64,
             input.shape(),
             device,
             instruction.input_grad_alloc_hints[0]
@@ -782,7 +798,7 @@ void batch_reduce_sum_mean_backward_impl(
         // input_grad.zero(device);
     }
 
-    Tensor grad(DataType::dt_float, {1}, device, AllocHint::temporary);
+    Tensor grad(DataType::dt_float64, {1}, device, AllocHint::temporary);
     if (keepdim) {
         auto output_grad =
             local_grads[instruction.output_indices[0]].contiguous(device);
@@ -799,7 +815,7 @@ void batch_reduce_sum_mean_backward_impl(
             )
         );
         Tensor temp(
-            DataType::dt_float,
+            DataType::dt_float64,
             {(temp_storage_bytes + 7) / 8},
             device,
             AllocHint::temporary
@@ -912,7 +928,7 @@ void op_offset_indices(
     std::size_t total_size = std::accumulate(sizes_out.begin(), sizes_out.end(), 0);
     auto& output = locals[instruction.output_indices[0]];
     output = Tensor(
-        DataType::dt_int, {total_size}, device, instruction.output_alloc_hints[0]
+        DataType::dt_int32, {total_size}, device, instruction.output_alloc_hints[0]
     );
     std::size_t sum_offset = 0, sum_out = 0;
     for (auto [size_offset, size_out] : zip(sizes_offset, sizes_out)) {
@@ -961,7 +977,7 @@ void op_quantile(
         return;
     }
 
-    Tensor tmp(DataType::dt_float, {batch_size}, device, AllocHint::temporary);
+    Tensor tmp(DataType::dt_float64, {batch_size}, device, AllocHint::temporary);
     tmp.copy_from(input, device);
     std::size_t temp_storage_bytes;
     check_error(
@@ -975,7 +991,10 @@ void op_quantile(
         )
     );
     Tensor tmp_sort(
-        DataType::dt_float, {(temp_storage_bytes + 7) / 8}, device, AllocHint::temporary
+        DataType::dt_float64,
+        {(temp_storage_bytes + 7) / 8},
+        device,
+        AllocHint::temporary
     );
     check_error(
         cub::DeviceMergeSort::SortKeys(
@@ -989,7 +1008,8 @@ void op_quantile(
     );
     tmp_sort.reset(device);
 
-    output = Tensor(DataType::dt_float, {1}, device, instruction.output_alloc_hints[0]);
+    output =
+        Tensor(DataType::dt_float64, {1}, device, instruction.output_alloc_hints[0]);
     launch_kernel(
         kernel_gather_quantile,
         1,
@@ -1022,7 +1042,10 @@ void op_random(
     auto& output = locals[instruction.output_indices[0]];
     auto dim = instruction.output_shapes[0][0];
     output = Tensor(
-        DataType::dt_float, {batch_size, dim}, device, instruction.output_alloc_hints[0]
+        DataType::dt_float64,
+        {batch_size, dim},
+        device,
+        instruction.output_alloc_hints[0]
     );
     launch_rng_kernel(
         kernel_random,
@@ -1058,7 +1081,7 @@ void op_random_int(
     auto max_val = locals[instruction.input_indices[1]].batch_sizes()[0];
     auto& output = locals[instruction.output_indices[0]];
     output = Tensor(
-        DataType::dt_int, {batch_size}, device, instruction.output_alloc_hints[0]
+        DataType::dt_int32, {batch_size}, device, instruction.output_alloc_hints[0]
     );
     launch_rng_kernel(
         kernel_random_int,
@@ -1106,16 +1129,18 @@ void op_unweight(
     gpuStream_t stream = device.stream();
 
     if (batch_size == 0) {
-        indices =
-            Tensor(DataType::dt_float, {0}, device, instruction.output_alloc_hints[0]);
-        uw_weights =
-            Tensor(DataType::dt_float, {0}, device, instruction.output_alloc_hints[1]);
+        indices = Tensor(
+            DataType::dt_float64, {0}, device, instruction.output_alloc_hints[0]
+        );
+        uw_weights = Tensor(
+            DataType::dt_float64, {0}, device, instruction.output_alloc_hints[1]
+        );
         return;
     }
 
-    Tensor indices_tmp(DataType::dt_int, {batch_size}, device, AllocHint::temporary);
+    Tensor indices_tmp(DataType::dt_int32, {batch_size}, device, AllocHint::temporary);
     Tensor uw_weights_tmp(
-        DataType::dt_float, {batch_size}, device, AllocHint::temporary
+        DataType::dt_float64, {batch_size}, device, AllocHint::temporary
     );
     launch_rng_kernel(
         kernel_unweight,
@@ -1129,7 +1154,7 @@ void op_unweight(
     );
 
     Tensor indices_compacted(
-        DataType::dt_int, {batch_size}, device, instruction.output_alloc_hints[0]
+        DataType::dt_int32, {batch_size}, device, instruction.output_alloc_hints[0]
     );
     auto ptr_all =
         thrust::device_pointer_cast(static_cast<me_int_t*>(indices_tmp.data()));
@@ -1145,8 +1170,9 @@ void op_unweight(
 
     std::size_t count = ptr_compacted_end - ptr_compacted;
     indices = indices_compacted.slice(0, 0, count);
-    uw_weights =
-        Tensor(DataType::dt_float, {count}, device, instruction.output_alloc_hints[1]);
+    uw_weights = Tensor(
+        DataType::dt_float64, {count}, device, instruction.output_alloc_hints[1]
+    );
     auto ptr_all_weights =
         thrust::device_pointer_cast(static_cast<double*>(uw_weights_tmp.data()));
     auto ptr_uw_weights =
@@ -1179,7 +1205,7 @@ void histogram_common(
 ) {
     auto policy = thrust_par.on(device.stream());
     Tensor reduce_tmp(
-        DataType::dt_float, {n_dims * n_bins}, device, AllocHint::temporary
+        DataType::dt_float64, {n_dims * n_bins}, device, AllocHint::temporary
     );
     auto indices_ptr =
         thrust::device_pointer_cast(static_cast<me_int_t*>(indices_tmp.data()));
@@ -1252,8 +1278,9 @@ void op_vegas_histogram(
     shape[0] = 1;
     std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1);
     values =
-        Tensor(DataType::dt_float, shape, device, instruction.output_alloc_hints[0]);
-    counts = Tensor(DataType::dt_int, shape, device, instruction.output_alloc_hints[1]);
+        Tensor(DataType::dt_float64, shape, device, instruction.output_alloc_hints[0]);
+    counts =
+        Tensor(DataType::dt_int32, shape, device, instruction.output_alloc_hints[1]);
 
     std::size_t batch_size = locals[instruction.batch_size_index].size(0);
     std::size_t n_dims = input.size(1);
@@ -1261,10 +1288,10 @@ void op_vegas_histogram(
 
     std::size_t padded_size = batch_size + n_bins;
     Tensor indices_tmp(
-        DataType::dt_int, {padded_size, n_dims}, device, AllocHint::temporary
+        DataType::dt_int32, {padded_size, n_dims}, device, AllocHint::temporary
     );
     Tensor weights_tmp(
-        DataType::dt_float, {padded_size, n_dims}, device, AllocHint::temporary
+        DataType::dt_float64, {padded_size, n_dims}, device, AllocHint::temporary
     );
     launch_kernel(
         kernel_prepare_vegas_hist,
@@ -1315,15 +1342,18 @@ void op_discrete_histogram(
     shape[0] = 1;
     std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1);
     values =
-        Tensor(DataType::dt_float, shape, device, instruction.output_alloc_hints[0]);
-    counts = Tensor(DataType::dt_int, shape, device, instruction.output_alloc_hints[1]);
+        Tensor(DataType::dt_float64, shape, device, instruction.output_alloc_hints[0]);
+    counts =
+        Tensor(DataType::dt_int32, shape, device, instruction.output_alloc_hints[1]);
 
     std::size_t batch_size = locals[instruction.batch_size_index].size(0);
     std::size_t n_opts = values.size(1);
 
     std::size_t padded_size = batch_size + n_opts;
-    Tensor indices_tmp(DataType::dt_int, {padded_size}, device, AllocHint::temporary);
-    Tensor weights_tmp(DataType::dt_float, {padded_size}, device, AllocHint::temporary);
+    Tensor indices_tmp(DataType::dt_int32, {padded_size}, device, AllocHint::temporary);
+    Tensor weights_tmp(
+        DataType::dt_float64, {padded_size}, device, AllocHint::temporary
+    );
     launch_kernel(
         kernel_prepare_discrete_hist,
         padded_size,
@@ -1396,21 +1426,23 @@ void op_histogram(
     shape[0] = 1;
     std::copy(out_shape.begin(), out_shape.end(), shape.begin() + 1);
     values =
-        Tensor(DataType::dt_float, shape, device, instruction.output_alloc_hints[0]);
+        Tensor(DataType::dt_float64, shape, device, instruction.output_alloc_hints[0]);
     square_values =
-        Tensor(DataType::dt_float, shape, device, instruction.output_alloc_hints[1]);
+        Tensor(DataType::dt_float64, shape, device, instruction.output_alloc_hints[1]);
 
     std::size_t batch_size = locals[instruction.batch_size_index].size(0);
     std::size_t n_bins = values.size(1) - 2;
     std::size_t padded_size = batch_size + n_bins + 2;
-    Tensor indices_tmp(DataType::dt_int, {padded_size}, device, AllocHint::temporary);
-    Tensor weights_tmp(DataType::dt_float, {padded_size}, device, AllocHint::temporary);
+    Tensor indices_tmp(DataType::dt_int32, {padded_size}, device, AllocHint::temporary);
+    Tensor weights_tmp(
+        DataType::dt_float64, {padded_size}, device, AllocHint::temporary
+    );
     Tensor square_weights_tmp(
-        DataType::dt_float, {padded_size}, device, AllocHint::temporary
+        DataType::dt_float64, {padded_size}, device, AllocHint::temporary
     );
     // the keys reduce_by_key writes: the n_bins bins, the underflow and the overflow
     // bin
-    Tensor reduce_tmp(DataType::dt_float, {n_bins + 2}, device, AllocHint::temporary);
+    Tensor reduce_tmp(DataType::dt_float64, {n_bins + 2}, device, AllocHint::temporary);
 
     launch_kernel(
         kernel_prepare_hist,
@@ -1794,6 +1826,13 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) :
         if (context->global_requires_grad(name)) {
             _requires_grad_init.at(value.local_index) = true;
             _grad_global_indices.push_back(value.local_index);
+            // the gradients of all globals are stored in one contiguous tensor
+            if (_grad_global_total_size > 0 && global.dtype() != _grad_global_dtype) {
+                throw std::invalid_argument(
+                    "all globals that require a gradient must have the same dtype"
+                );
+            }
+            _grad_global_dtype = global.dtype();
             _grad_global_shapes.push_back(global.shape());
             _grad_global_total_size += global.shape().product();
         }
@@ -2141,7 +2180,7 @@ std::pair<TensorVec, TensorVec> GpuRuntime::run_backward(
         }
     }
     Tensor all_global_grads(
-        DataType::dt_float,
+        _grad_global_dtype,
         {_grad_global_total_size},
         init_device,
         AllocHint::global_grad
@@ -2168,9 +2207,9 @@ std::pair<TensorVec, TensorVec> GpuRuntime::run_backward(
             for (auto [output_index, output_dtype] :
                  zip(instr.output_indices, instr.output_dtypes)) {
                 auto& grad = local_grads[output_index];
-                if (!grad && output_dtype == DataType::dt_float) {
+                if (!grad && is_float(output_dtype)) {
                     grad = Tensor(
-                        DataType::dt_float,
+                        output_dtype,
                         locals[output_index].shape(),
                         device,
                         AllocHint::local_grad
diff --git a/madspace/src/gpu/runtime.hpp b/madspace/src/gpu/runtime.hpp
index 212e1f479a..15b85a0880 100644
--- a/madspace/src/gpu/runtime.hpp
+++ b/madspace/src/gpu/runtime.hpp
@@ -91,6 +91,7 @@ class GpuRuntime : public Runtime {
     SizeVec _grad_global_indices;
     std::vector<Sizes> _grad_global_shapes;
     std::size_t _grad_global_total_size;
+    DataType _grad_global_dtype = DataType::dt_float64;
     ContextPtr _context;
     ThreadResource<std::vector<gpuStream_t>> _streams;
     ThreadResource<std::vector<gpuEvent_t>> _events;
diff --git a/madspace/src/kernels/nn.hpp b/madspace/src/kernels/nn.hpp
index 176329a5e8..4cef7f7067 100644
--- a/madspace/src/kernels/nn.hpp
+++ b/madspace/src/kernels/nn.hpp
@@ -8,102 +8,153 @@ namespace kernels {
 
 // Kernels
 
+// The activations are instantiated in float32 and float64. Literals are cast to
+// FVal<T> so that float32 code does not get promoted to double.
+
 template <typename T>
 KERNELSPEC void kernel_relu(FIn<T, 0> input, FOut<T, 0> output) {
-    FVal<T> x(input);
-    output = where(x < 0., 0., x);
+    using F = FVal<T>;
+    F x(input);
+    output = where(x < F(0.), F(0.), x);
 }
 
 template <typename T>
 KERNELSPEC void
 backward_kernel_relu(FIn<T, 0> input, FIn<T, 0> output_grad, FOut<T, 0> input_grad) {
-    FVal<T> x(input), g(output_grad);
-    input_grad = where(x < 0., 0., g);
+    using F = FVal<T>;
+    F x(input), g(output_grad);
+    input_grad += where(x < F(0.), F(0.), g);
 }
 
 template <typename T>
 KERNELSPEC void kernel_leaky_relu(FIn<T, 0> input, FOut<T, 0> output) {
-    FVal<T> x(input);
-    output = where(x < 0., x * 0.01, x);
+    using F = FVal<T>;
+    F x(input);
+    output = where(x < F(0.), x * F(0.01), x);
 }
 
 template <typename T>
 KERNELSPEC void backward_kernel_leaky_relu(
     FIn<T, 0> input, FIn<T, 0> output_grad, FOut<T, 0> input_grad
 ) {
-    FVal<T> x(input), g(output_grad);
-    input_grad = where(x < 0., g * 0.01, g);
+    using F = FVal<T>;
+    F x(input), g(output_grad);
+    input_grad += where(x < F(0.), g * F(0.01), g);
 }
 
 template <typename T>
 KERNELSPEC void kernel_elu(FIn<T, 0> input, FOut<T, 0> output) {
-    FVal<T> x(input);
-    output = where(x < 0., expm1(x), x);
+    using F = FVal<T>;
+    F x(input);
+    output = where(x < F(0.), expm1(x), x);
 }
 
 template <typename T>
 KERNELSPEC void
 backward_kernel_elu(FIn<T, 0> input, FIn<T, 0> output_grad, FOut<T, 0> input_grad) {
-    FVal<T> x(input), g(output_grad);
-    input_grad = where(x < 0., g * exp(x), g);
+    using F = FVal<T>;
+    F x(input), g(output_grad);
+    input_grad += where(x < F(0.), g * exp(x), g);
 }
 
 template <typename T>
 KERNELSPEC void kernel_gelu(FIn<T, 0> input, FOut<T, 0> output) {
-    output = 0.5 * input * (1. + erf(SQRT_HALF * input));
+    using F = FVal<T>;
+    F x(input);
+    output = F(0.5) * x * (F(1.) + erf(F(SQRT_HALF) * x));
 }
 
 template <typename T>
 KERNELSPEC void
 backward_kernel_gelu(FIn<T, 0> input, FIn<T, 0> output_grad, FOut<T, 0> input_grad) {
-    FVal<T> x(input), g(output_grad);
-    auto cdf = 0.5 * (1. + erf(SQRT_HALF * x));
-    auto pdf = 0.5 * SQRT_HALF * TWO_DIV_SQRT_PI * exp(-0.5 * x * x);
-    input_grad = g * (cdf + x * pdf);
+    using F = FVal<T>;
+    F x(input), g(output_grad);
+    F cdf = F(0.5) * (F(1.) + erf(F(SQRT_HALF) * x));
+    F pdf = F(0.5 * SQRT_HALF * TWO_DIV_SQRT_PI) * exp(F(-0.5) * x * x);
+    input_grad += g * (cdf + x * pdf);
 }
 
 template <typename T>
 KERNELSPEC void kernel_sigmoid(FIn<T, 0> input, FOut<T, 0> output) {
-    FVal<T> x = input;
-    output = 1. / (1. + exp(-x));
+    using F = FVal<T>;
+    F x(input);
+    output = F(1.) / (F(1.) + exp(-x));
 }
 
 template <typename T>
 KERNELSPEC void backward_kernel_sigmoid(
     FIn<T, 0> output, FIn<T, 0> output_grad, FOut<T, 0> input_grad
 ) {
-    FVal<T> y(output), g(output_grad);
-    input_grad = g * y * (1. - y);
+    using F = FVal<T>;
+    F y(output), g(output_grad);
+    input_grad += g * y * (F(1.) - y);
 }
 
 template <typename T>
 KERNELSPEC void kernel_softplus(FIn<T, 0> input, FOut<T, 0> output) {
-    FVal<T> x(input);
-    output = where(x < 20., log1p(exp(x)), x);
+    using F = FVal<T>;
+    F x(input);
+    output = where(x < F(20.), log1p(exp(x)), x);
 }
 
 template <typename T>
 KERNELSPEC void backward_kernel_softplus(
     FIn<T, 0> input, FIn<T, 0> output_grad, FOut<T, 0> input_grad
 ) {
-    FVal<T> x(input), g(output_grad);
-    auto z = exp(x);
-    input_grad = where(x < 20., g * z / (1. + z), g);
+    using F = FVal<T>;
+    F x(input), g(output_grad);
+    F z = exp(x);
+    input_grad += where(x < F(20.), g * z / (F(1.) + z), g);
 }
 
 template <typename T>
 KERNELSPEC void kernel_scaled_tanh(FIn<T, 0> input, FOut<T, 0> output) {
-    FVal<T> x(input);
-    output = 4. * tanh(x / 4.);
+    using F = FVal<T>;
+    F x(input);
+    output = F(4.) * tanh(x / F(4.));
 }
 
 template <typename T>
 KERNELSPEC void backward_kernel_scaled_tanh(
     FIn<T, 0> output, FIn<T, 0> output_grad, FOut<T, 0> input_grad
 ) {
-    FVal<T> g(output_grad);
-    FVal<T> y_scaled = output / 4.;
-    input_grad = g * (1. - y_scaled) * (1. + y_scaled);
+    using F = FVal<T>;
+    F g(output_grad), y(output);
+    F y_scaled = y / F(4.);
+    input_grad += g * (F(1.) - y_scaled) * (F(1.) + y_scaled);
+}
+
+// Precision conversions. T is the policy of either precision; T::F32 and T::F64
+// select the two sides.
+
+template <typename T>
+KERNELSPEC void
+kernel_convert_float32(FIn<typename T::F64, 0> input, FOut<typename T::F32, 0> output) {
+    output = FVal<typename T::F32>(FVal<typename T::F64>(input));
+}
+
+template <typename T>
+KERNELSPEC void backward_kernel_convert_float32(
+    FIn<typename T::F64, 0> input,
+    FIn<typename T::F32, 0> output_grad,
+    FOut<typename T::F64, 0> input_grad
+) {
+    input_grad += FVal<typename T::F64>(FVal<typename T::F32>(output_grad));
+}
+
+template <typename T>
+KERNELSPEC void
+kernel_convert_float64(FIn<typename T::F32, 0> input, FOut<typename T::F64, 0> output) {
+    output = FVal<typename T::F64>(FVal<typename T::F32>(input));
+}
+
+template <typename T>
+KERNELSPEC void backward_kernel_convert_float64(
+    FIn<typename T::F32, 0> input,
+    FIn<typename T::F64, 0> output_grad,
+    FOut<typename T::F32, 0> input_grad
+) {
+    input_grad += FVal<typename T::F32>(FVal<typename T::F64>(output_grad));
 }
 
 template <typename T>
diff --git a/madspace/src/kernels/operations.hpp b/madspace/src/kernels/operations.hpp
index b129ef6e8f..1937c8ae50 100644
--- a/madspace/src/kernels/operations.hpp
+++ b/madspace/src/kernels/operations.hpp
@@ -81,13 +81,12 @@ void backward_batch_foreach(
             if (!input_grad || input_grad.size(0) != batch_size) {
                 Sizes shape = input.shape();
                 shape[0] = batch_size;
-                input_grad =
-                    Tensor(input.dtype(), shape, device, AllocHint::temporary);
+                input_grad = Tensor(input.dtype(), shape, device, AllocHint::temporary);
             }
         } else if (!input_grad) {
             auto& input = locals[input_index];
-            /*if (input.size(0) != batch_size && input.dtype() == DataType::dt_float) {
-                throw std::runtime_error("backward not possible for broadcasting args");
+            /*if (input.size(0) != batch_size && input.dtype() == DataType::dt_float64)
+            { throw std::runtime_error("backward not possible for broadcasting args");
             }*/
             input_grad = Tensor(
                 input.dtype(),
@@ -104,6 +103,11 @@ void backward_batch_foreach(
     foreach_func(args, input_grads, batch_size, device);
 }
 
+// gradients have the dtype of float inputs, integer inputs get float64 gradients
+inline DataType grad_dtype(const Tensor& input) {
+    return is_float(input.dtype()) ? input.dtype() : DataType::dt_float64;
+}
+
 template <typename I, typename D>
 void op_stack(const I& instruction, TensorVec& locals, const D& device) {
     auto& first_shape = locals[instruction.input_indices[0]].shape();
@@ -129,13 +133,12 @@ template <typename I, typename D>
 void backward_op_stack(
     const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device
 ) {
-    // TODO: differentiate integer and float here (also other backwards)
     auto& output_grad = local_grads[instruction.output_indices[0]];
     for (std::size_t i = 0; std::size_t input_index : instruction.input_indices) {
         auto& input_grad = local_grads[input_index];
         if (!input_grad) {
             input_grad = Tensor(
-                DataType::dt_float,
+                grad_dtype(locals[input_index]),
                 locals[input_index].shape(),
                 device,
                 instruction.input_grad_alloc_hints[i]
@@ -186,7 +189,7 @@ void backward_op_unstack(
     auto& input_grad = local_grads[input_index];
     if (!input_grad) {
         input_grad = Tensor(
-            DataType::dt_float,
+            grad_dtype(locals[input_index]),
             locals[input_index].shape(),
             device,
             instruction.input_grad_alloc_hints[0]
@@ -213,10 +216,10 @@ void backward_op_pop(
     const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device
 ) {
     auto input_index = instruction.input_indices[0];
-    auto input_grad = local_grads[input_index];
+    auto& input_grad = local_grads[input_index];
     if (!input_grad) {
         input_grad = Tensor(
-            DataType::dt_float,
+            grad_dtype(locals[input_index]),
             locals[input_index].shape(),
             device,
             instruction.input_grad_alloc_hints[0]
@@ -269,7 +272,7 @@ void backward_op_batch_cat(
         auto& input_grad = local_grads[input_index];
         if (!input_grad) {
             input_grad = Tensor(
-                DataType::dt_float,
+                grad_dtype(locals[input_index]),
                 locals[input_index].shape(),
                 device,
                 instruction.input_grad_alloc_hints[i]
@@ -305,7 +308,7 @@ void backward_op_batch_split(
     auto& input_grad = local_grads[input_index];
     if (!input_grad) {
         input_grad = Tensor(
-            DataType::dt_float,
+            grad_dtype(locals[input_index]),
             locals[input_index].shape(),
             device,
             instruction.input_grad_alloc_hints[0]
@@ -357,7 +360,7 @@ void backward_op_cat(
         auto& input_grad = local_grads[input_index];
         if (!input_grad) {
             input_grad = Tensor(
-                DataType::dt_float,
+                grad_dtype(locals[input_index]),
                 locals[input_index].shape(),
                 device,
                 instruction.input_grad_alloc_hints[i]
@@ -471,7 +474,7 @@ void backward_op_rqs_reshape(
     auto& input_grad = local_grads[input_index];
     if (!input_grad) {
         input_grad = Tensor(
-            DataType::dt_float,
+            grad_dtype(locals[input_index]),
             locals[input_index].shape(),
             device,
             instruction.input_grad_alloc_hints[0]
diff --git a/madspace/src/kernels/util.hpp b/madspace/src/kernels/util.hpp
index 5254ce0d7c..5a9cd90c8f 100644
--- a/madspace/src/kernels/util.hpp
+++ b/madspace/src/kernels/util.hpp
@@ -37,7 +37,7 @@ template <typename T>
 KERNELSPEC void backward_kernel_gather(
     IIn<T, 0> index,
     FIn<T, 0> output_grad,
-    FOut<T, 0> index_grad,
+    IOut<T, 0> index_grad,
     FOut<T, 1> choices_grad
 ) {
     choices_grad.scatter_add(index, output_grad);
@@ -78,7 +78,7 @@ KERNELSPEC void backward_kernel_select(
     ISIn<T, 1> indices,
     FIn<T, 1> output_grad,
     FOut<T, 1> input_grad,
-    FOut<T, 1> indices_grad
+    IOut<T, 1> indices_grad
 ) {
     for (std::size_t i = 0; i < indices.size(); ++i) {
         input_grad.scatter_add(indices[i], output_grad[i]);
diff --git a/madspace/src/phasespace/channel_weight_network.cpp b/madspace/src/phasespace/channel_weight_network.cpp
index 343468299f..f07e3d7cf9 100644
--- a/madspace/src/phasespace/channel_weight_network.cpp
+++ b/madspace/src/phasespace/channel_weight_network.cpp
@@ -33,7 +33,8 @@ ChannelWeightNetwork::ChannelWeightNetwork(
     std::size_t layers,
     MLP::Activation activation,
     const std::string& prefix,
-    bool include_preprocessing
+    bool include_preprocessing,
+    DataType precision
 ) :
     FunctionGenerator(
         "ChannelWeightNetwork",
@@ -48,7 +49,8 @@ ChannelWeightNetwork::ChannelWeightNetwork(
         hidden_dim,
         layers,
         activation,
-        prefix
+        prefix,
+        precision
     ),
     _channel_count(channel_count),
     _mask_name(prefixed_name(prefix, "active_channels_mask")) {}
@@ -57,7 +59,7 @@ NamedVector<Value> ChannelWeightNetwork::build_function_impl(
     FunctionBuilder& fb, const NamedVector<Value>& args
 ) const {
     auto mask =
-        fb.global(_mask_name, DataType::dt_float, {static_cast<int>(_channel_count)});
+        fb.global(_mask_name, DataType::dt_float64, {static_cast<int>(_channel_count)});
     // auto net_input =
     //     _preprocessing.build_function(fb, {args["momenta"], args["x1"], args["x2"]});
     // auto net_output = _mlp.build_function(fb, net_input.values()).at(0);
@@ -73,14 +75,14 @@ void ChannelWeightNetwork::initialize_globals(
 ) const {
     _mlp.initialize_globals(context, seed);
 
-    context->define_global(_mask_name, DataType::dt_float, {_channel_count});
+    context->define_global(_mask_name, DataType::dt_float64, {_channel_count});
     bool is_cpu = context->device() == cpu_device();
     auto mask_global = context->global(_mask_name);
     Tensor mask;
     if (is_cpu) {
         mask = mask_global;
     } else {
-        mask = Tensor(DataType::dt_float, mask_global.shape());
+        mask = Tensor(DataType::dt_float64, mask_global.shape());
     }
     auto mask_view = mask.view<double, 2>()[0];
     for (std::size_t i = 0; i < mask_view.size(); ++i) {
diff --git a/madspace/src/phasespace/discrete_flow.cpp b/madspace/src/phasespace/discrete_flow.cpp
index 5f283fcb19..bf5c899294 100644
--- a/madspace/src/phasespace/discrete_flow.cpp
+++ b/madspace/src/phasespace/discrete_flow.cpp
@@ -12,7 +12,8 @@ DiscreteFlow::DiscreteFlow(
     std::size_t condition_dim,
     std::size_t subnet_hidden_dim,
     std::size_t subnet_layers,
-    MLP::Activation subnet_activation
+    MLP::Activation subnet_activation,
+    DataType subnet_precision
 ) :
     Mapping(
         "DiscreteFlow",
@@ -57,7 +58,8 @@ DiscreteFlow::DiscreteFlow(
                 subnet_hidden_dim,
                 subnet_layers,
                 subnet_activation,
-                prefixed_name(prefix, std::format("subnet{}_", dim_index + 1))
+                prefixed_name(prefix, std::format("subnet{}_", dim_index + 1)),
+                subnet_precision
             );
         } else {
             _first_prob_name = prefixed_name(prefix, "prob0");
@@ -112,7 +114,7 @@ Mapping::Result DiscreteFlow::build_transform(
         if (dim_index == 0 && _first_prob_name) {
             probs = fb.global(
                 _first_prob_name.value(),
-                DataType::dt_float,
+                DataType::dt_float64,
                 {static_cast<int>(option_count)}
             );
         } else {
diff --git a/madspace/src/phasespace/discrete_sampler.cpp b/madspace/src/phasespace/discrete_sampler.cpp
index 39c6db5cf7..26c475cb26 100644
--- a/madspace/src/phasespace/discrete_sampler.cpp
+++ b/madspace/src/phasespace/discrete_sampler.cpp
@@ -111,8 +111,9 @@ Mapping::Result DiscreteSampler::build_transform(
     std::size_t condition_index = 0;
     for (auto [input, prob_name, option_count, has_prior] :
          zip(inputs, _prob_names, _option_counts, _dim_has_prior)) {
-        auto probs =
-            fb.global(prob_name, DataType::dt_float, {static_cast<int>(option_count)});
+        auto probs = fb.global(
+            prob_name, DataType::dt_float64, {static_cast<int>(option_count)}
+        );
         if (has_prior) {
             probs = fb.mul(probs, conditions.at(condition_index));
             ++condition_index;
@@ -139,8 +140,9 @@ void madspace::initialize_uniform_probs(
     ContextPtr context, const std::string& name, std::size_t option_count
 ) {
     bool is_cpu = context->device() == cpu_device();
-    auto prob_global = context->define_global(name, DataType::dt_float, {option_count});
-    auto prob = is_cpu ? prob_global : Tensor(DataType::dt_float, {1, option_count});
+    auto prob_global =
+        context->define_global(name, DataType::dt_float64, {option_count});
+    auto prob = is_cpu ? prob_global : Tensor(DataType::dt_float64, {1, option_count});
     auto prob_view = prob.view<double, 2>()[0];
     auto prob_value = 1. / option_count;
     for (std::size_t i = 0; i < prob_view.size(); ++i) {
diff --git a/madspace/src/phasespace/flow.cpp b/madspace/src/phasespace/flow.cpp
index 1f353fcf1e..ae7a15060f 100644
--- a/madspace/src/phasespace/flow.cpp
+++ b/madspace/src/phasespace/flow.cpp
@@ -46,7 +46,7 @@ void vegas_init(
     auto grid_tensor = context->global(grid_name).cpu();
     bool is_cpu = false; // context->device() == cpu_device();
     Tensor bias_tensor =
-        is_cpu ? bias_global : Tensor(DataType::dt_float, bias_global.shape());
+        is_cpu ? bias_global : Tensor(DataType::dt_float64, bias_global.shape());
     auto bias_view = bias_tensor.view<double, 2>()[0];
     auto grid_view = grid_tensor.view<double, 3>()[0];
     std::size_t n_dims = grid_tensor.size(1);
@@ -142,7 +142,7 @@ void vegas_init(
     }
 
     if (!is_cpu) {
-        bias_global.copy_from(bias_tensor);
+        bias_global.copy_from(bias_tensor.cast(bias_global.dtype()));
     }
 }
 
@@ -156,7 +156,8 @@ Flow::Flow(
     std::size_t subnet_hidden_dim,
     std::size_t subnet_layers,
     MLP::Activation subnet_activation,
-    bool invert_spline
+    bool invert_spline,
+    DataType subnet_precision
 ) :
     Mapping(
         "Flow",
@@ -198,13 +199,15 @@ Flow::Flow(
                  subnet_hidden_dim,
                  subnet_layers,
                  subnet_activation,
-                 prefixed_name(prefix, std::format("subnet{}a", block_index + 1))),
+                 prefixed_name(prefix, std::format("subnet{}a", block_index + 1)),
+                 subnet_precision),
              MLP(indices1.size() + condition_dim,
                  indices2.size() * (3 * bin_count + 1),
                  subnet_hidden_dim,
                  subnet_layers,
                  subnet_activation,
-                 prefixed_name(prefix, std::format("subnet{}b", block_index + 1))),
+                 prefixed_name(prefix, std::format("subnet{}b", block_index + 1)),
+                 subnet_precision),
              indices1,
              indices2}
         );
diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp
index 6cb68cc243..8dd83a6f1b 100644
--- a/madspace/src/phasespace/integrand.cpp
+++ b/madspace/src/phasespace/integrand.cpp
@@ -291,13 +291,13 @@ NamedVector<Value> Integrand::build_function_impl(
 }
 
 NamedVector<Type> Integrand::compute_channel_part_ret_types() const {
-    auto acc_float = Type(DataType::dt_float, acc_batch_size, {});
-    auto acc_int = Type(DataType::dt_int, acc_batch_size, {});
+    auto acc_float = Type(DataType::dt_float64, acc_batch_size, {});
+    auto acc_int = Type(DataType::dt_int32, acc_batch_size, {});
     auto acc_float_array = [](int n) {
-        return Type(DataType::dt_float, acc_batch_size, {n});
+        return Type(DataType::dt_float64, acc_batch_size, {n});
     };
     auto acc_four_vec_array = [](int n) {
-        return Type(DataType::dt_float, acc_batch_size, {n, 4});
+        return Type(DataType::dt_float64, acc_batch_size, {n, 4});
     };
 
     bool has_multi_flavor = _pid_options.size() > 1;
@@ -321,7 +321,7 @@ NamedVector<Type> Integrand::compute_channel_part_ret_types() const {
     }
 
     // outputs after cuts
-    ret.push_back("indices_acc", Type(DataType::dt_int, acc_batch_size, {}));
+    ret.push_back("indices_acc", Type(DataType::dt_int32, acc_batch_size, {}));
     ret.push_back("momenta_acc", acc_four_vec_array(particle_count));
     ret.push_back("x1_acc", acc_float);
     ret.push_back("x2_acc", acc_float);
diff --git a/madspace/src/phasespace/mlp.cpp b/madspace/src/phasespace/mlp.cpp
index 4de6a6537f..8da85ee264 100644
--- a/madspace/src/phasespace/mlp.cpp
+++ b/madspace/src/phasespace/mlp.cpp
@@ -15,16 +15,17 @@ Value build_layer(
     int output_dim,
     MLP::Activation activation,
     const std::string& prefix,
-    int layer_index
+    int layer_index,
+    DataType precision
 ) {
     auto weight = fb.global(
         prefixed_name(prefix, std::format("layer{}.weight", layer_index)),
-        DataType::dt_float,
+        precision,
         {output_dim, input_dim}
     );
     auto bias = fb.global(
         prefixed_name(prefix, std::format("layer{}.bias", layer_index)),
-        DataType::dt_float,
+        precision,
         {output_dim}
     );
     auto linear_out = fb.matmul(input, weight, bias);
@@ -53,25 +54,19 @@ void initialize_layer(
     const std::string& prefix,
     int layer_index,
     std::optional<std::uint64_t> seed,
-    bool zeros
+    bool zeros,
+    DataType precision
 ) {
     auto weight_name =
         prefixed_name(prefix, std::format("layer{}.weight", layer_index));
     auto bias_name = prefixed_name(prefix, std::format("layer{}.bias", layer_index));
-    auto weight_tensor_global = context->define_global(
-        weight_name, DataType::dt_float, {output_dim, input_dim}, true
-    );
+    auto weight_tensor_global =
+        context->define_global(weight_name, precision, {output_dim, input_dim}, true);
     auto bias_tensor_global =
-        context->define_global(bias_name, DataType::dt_float, {output_dim}, true);
-    Tensor weight_tensor, bias_tensor;
-    bool is_cpu = context->device() == cpu_device();
-    if (is_cpu) {
-        weight_tensor = weight_tensor_global;
-        bias_tensor = bias_tensor_global;
-    } else {
-        weight_tensor = Tensor(DataType::dt_float, {1, output_dim, input_dim});
-        bias_tensor = Tensor(DataType::dt_float, {1, output_dim});
-    }
+        context->define_global(bias_name, precision, {output_dim}, true);
+    // filled in double precision on the CPU, then converted and copied
+    Tensor weight_tensor(DataType::dt_float64, {1, output_dim, input_dim});
+    Tensor bias_tensor(DataType::dt_float64, {1, output_dim});
 
     auto weight_view = weight_tensor.view<double, 3>()[0];
     auto bias_view = bias_tensor.view<double, 2>()[0];
@@ -105,10 +100,8 @@ void initialize_layer(
         }
     }
 
-    if (!is_cpu) {
-        weight_tensor_global.copy_from(weight_tensor);
-        bias_tensor_global.copy_from(bias_tensor);
-    }
+    weight_tensor_global.copy_from(weight_tensor.cast(precision));
+    bias_tensor_global.copy_from(bias_tensor.cast(precision));
 }
 
 } // namespace
@@ -119,7 +112,8 @@ MLP::MLP(
     std::size_t hidden_dim,
     std::size_t layers,
     Activation activation,
-    const std::string& prefix
+    const std::string& prefix,
+    DataType precision
 ) :
     FunctionGenerator(
         "MLP",
@@ -131,7 +125,11 @@ MLP::MLP(
     _hidden_dim(hidden_dim),
     _layers(layers),
     _activation(activation),
-    _prefix(prefix) {
+    _prefix(prefix),
+    _precision(precision) {
+    if (!is_float(precision)) {
+        throw std::invalid_argument("MLP precision must be a float type");
+    }
     if (input_dim == 0) {
         throw std::invalid_argument("MLP input dimension cannot be 0");
     }
@@ -142,15 +140,16 @@ MLP::MLP(
 
 NamedVector<Value>
 MLP::build_function_impl(FunctionBuilder& fb, const NamedVector<Value>& args) const {
+    bool single = _precision == DataType::dt_float32;
     std::size_t dim = _input_dim;
-    Value x = args.at(0);
+    Value x = single ? fb.convert_float32(args.at(0)) : args.at(0);
     for (std::size_t i = 1; i < _layers; ++i) {
-        x = build_layer(fb, x, dim, _hidden_dim, _activation, _prefix, i);
+        x = build_layer(fb, x, dim, _hidden_dim, _activation, _prefix, i, _precision);
         dim = _hidden_dim;
     }
-    return {
-        {"output", build_layer(fb, x, dim, _output_dim, MLP::linear, _prefix, _layers)}
-    };
+    Value output =
+        build_layer(fb, x, dim, _output_dim, MLP::linear, _prefix, _layers, _precision);
+    return {{"output", single ? fb.convert_float64(output) : output}};
 }
 
 void MLP::initialize_globals(
@@ -158,10 +157,14 @@ void MLP::initialize_globals(
 ) const {
     std::size_t dim = _input_dim;
     for (std::size_t i = 1; i < _layers; ++i) {
-        initialize_layer(context, dim, _hidden_dim, _prefix, i, seed, false);
+        initialize_layer(
+            context, dim, _hidden_dim, _prefix, i, seed, false, _precision
+        );
         dim = _hidden_dim;
     }
-    initialize_layer(context, dim, _output_dim, _prefix, _layers, seed, true);
+    initialize_layer(
+        context, dim, _output_dim, _prefix, _layers, seed, true, _precision
+    );
 }
 
 std::vector<std::string> MLP::global_names() const {
diff --git a/madspace/src/phasespace/pdf.cpp b/madspace/src/phasespace/pdf.cpp
index 0a437f341e..63bdc3f8e3 100644
--- a/madspace/src/phasespace/pdf.cpp
+++ b/madspace/src/phasespace/pdf.cpp
@@ -320,14 +320,14 @@ std::vector<std::size_t> PdfGrid::logq2_shape(bool batch_dim) const {
 
 void PdfGrid::initialize_globals(ContextPtr context, const std::string& prefix) const {
     auto logx_tensor_global = context->define_global(
-        prefixed_name(prefix, "pdf_logx"), DataType::dt_float, logx_shape()
+        prefixed_name(prefix, "pdf_logx"), DataType::dt_float64, logx_shape()
     );
     auto q2_tensor_global = context->define_global(
-        prefixed_name(prefix, "pdf_logq2"), DataType::dt_float, logq2_shape()
+        prefixed_name(prefix, "pdf_logq2"), DataType::dt_float64, logq2_shape()
     );
     auto coeffs_tensor_global = context->define_global(
         prefixed_name(prefix, "pdf_coefficients"),
-        DataType::dt_float,
+        DataType::dt_float64,
         coefficients_shape()
     );
     bool is_cpu = context->device() == cpu_device();
@@ -337,9 +337,9 @@ void PdfGrid::initialize_globals(ContextPtr context, const std::string& prefix)
         q2_tensor = q2_tensor_global;
         coeffs_tensor = coeffs_tensor_global;
     } else {
-        logx_tensor = Tensor(DataType::dt_float, logx_shape(true));
-        q2_tensor = Tensor(DataType::dt_float, logq2_shape(true));
-        coeffs_tensor = Tensor(DataType::dt_float, coefficients_shape(true));
+        logx_tensor = Tensor(DataType::dt_float64, logx_shape(true));
+        q2_tensor = Tensor(DataType::dt_float64, logq2_shape(true));
+        coeffs_tensor = Tensor(DataType::dt_float64, coefficients_shape(true));
     }
     initialize_logx(logx_tensor);
     initialize_logq2(q2_tensor);
@@ -391,17 +391,17 @@ NamedVector<Value> PartonDensity::build_function_impl(
     auto q = args.at(1);
     auto grid_logx = fb.global(
         prefixed_name(_prefix, "pdf_logx"),
-        DataType::dt_float,
+        DataType::dt_float64,
         {_logx_shape.begin(), _logx_shape.end()}
     );
     auto grid_logq2 = fb.global(
         prefixed_name(_prefix, "pdf_logq2"),
-        DataType::dt_float,
+        DataType::dt_float64,
         {_logq2_shape.begin(), _logq2_shape.end()}
     );
     auto grid_coeffs = fb.global(
         prefixed_name(_prefix, "pdf_coefficients"),
-        DataType::dt_float,
+        DataType::dt_float64,
         {_coeffs_shape.begin(), _coeffs_shape.end()}
     );
     if (_dynamic_pid) {
@@ -546,11 +546,11 @@ void AlphaSGrid::initialize_globals(
     ContextPtr context, const std::string& prefix
 ) const {
     auto q2_tensor_global = context->define_global(
-        prefixed_name(prefix, "alpha_s_logq2"), DataType::dt_float, logq2_shape()
+        prefixed_name(prefix, "alpha_s_logq2"), DataType::dt_float64, logq2_shape()
     );
     auto coeffs_tensor_global = context->define_global(
         prefixed_name(prefix, "alpha_s_coefficients"),
-        DataType::dt_float,
+        DataType::dt_float64,
         coefficients_shape()
     );
     bool is_cpu = context->device() == cpu_device();
@@ -559,8 +559,8 @@ void AlphaSGrid::initialize_globals(
         q2_tensor = q2_tensor_global;
         coeffs_tensor = coeffs_tensor_global;
     } else {
-        q2_tensor = Tensor(DataType::dt_float, logq2_shape(true));
-        coeffs_tensor = Tensor(DataType::dt_float, coefficients_shape(true));
+        q2_tensor = Tensor(DataType::dt_float64, logq2_shape(true));
+        coeffs_tensor = Tensor(DataType::dt_float64, coefficients_shape(true));
     }
     initialize_logq2(q2_tensor);
     initialize_coefficients(coeffs_tensor);
@@ -584,12 +584,12 @@ NamedVector<Value> RunningCoupling::build_function_impl(
     auto q = args.at(0);
     auto grid_logq2 = fb.global(
         prefixed_name(_prefix, "alpha_s_logq2"),
-        DataType::dt_float,
+        DataType::dt_float64,
         {_logq2_shape.begin(), _logq2_shape.end()}
     );
     auto grid_coeffs = fb.global(
         prefixed_name(_prefix, "alpha_s_coefficients"),
-        DataType::dt_float,
+        DataType::dt_float64,
         {_coeffs_shape.begin(), _coeffs_shape.end()}
     );
     return {{"alpha_s", fb.interpolate_alpha_s(q, grid_logq2, grid_coeffs)}};
diff --git a/madspace/src/phasespace/phasespace.cpp b/madspace/src/phasespace/phasespace.cpp
index 11fff4c9d3..d1f257d77c 100644
--- a/madspace/src/phasespace/phasespace.cpp
+++ b/madspace/src/phasespace/phasespace.cpp
@@ -276,7 +276,7 @@ PhaseSpaceMapping::PhaseSpaceMapping(
             if (nd > 0) {
                 in.push_back(
                     "discrete",
-                    Type{DataType::dt_int, batch_size, {static_cast<int>(nd)}}
+                    Type{DataType::dt_int32, batch_size, {static_cast<int>(nd)}}
                 );
             }
             return in;
diff --git a/madspace/src/phasespace/vegas.cpp b/madspace/src/phasespace/vegas.cpp
index 0d5f804cb4..4f3e03e456 100644
--- a/madspace/src/phasespace/vegas.cpp
+++ b/madspace/src/phasespace/vegas.cpp
@@ -39,7 +39,7 @@ Mapping::Result VegasMapping::build_forward_impl(
 ) const {
     auto grid = fb.global(
         _grid_name,
-        DataType::dt_float,
+        DataType::dt_float64,
         {static_cast<int>(_dimension), static_cast<int>(_bin_count) + 1}
     );
     auto [output, dets] = fb.vegas_forward(inputs.at(0), grid);
@@ -53,7 +53,7 @@ Mapping::Result VegasMapping::build_inverse_impl(
 ) const {
     auto grid = fb.global(
         _grid_name,
-        DataType::dt_float,
+        DataType::dt_float64,
         {static_cast<int>(_dimension), static_cast<int>(_bin_count) + 1}
     );
     auto [output, dets] = fb.vegas_inverse(inputs.at(0), grid);
@@ -62,7 +62,7 @@ Mapping::Result VegasMapping::build_inverse_impl(
 
 void VegasMapping::initialize_globals(ContextPtr context) const {
     context->define_global(
-        _grid_name, DataType::dt_float, {_dimension, _bin_count + 1}
+        _grid_name, DataType::dt_float64, {_dimension, _bin_count + 1}
     );
     initialize_vegas_grid(context, _grid_name);
 }
@@ -71,14 +71,14 @@ void madspace::initialize_vegas_grid(ContextPtr context, const std::string& grid
     bool is_cpu = context->device() == cpu_device();
     auto grid_global = context->global(grid_name);
     if (grid_global.shape().size() != 3 || grid_global.size(0) != 1 ||
-        grid_global.size(2) < 2 || grid_global.dtype() != DataType::dt_float) {
+        grid_global.size(2) < 2 || grid_global.dtype() != DataType::dt_float64) {
         throw std::runtime_error("Invalid vegas grid type");
     }
     Tensor grid;
     if (is_cpu) {
         grid = grid_global;
     } else {
-        grid = Tensor(DataType::dt_float, grid_global.shape());
+        grid = Tensor(DataType::dt_float64, grid_global.shape());
     }
 
     auto grid_view = grid.view<double, 3>()[0];
diff --git a/madspace/src/python/function_runtime.cpp b/madspace/src/python/function_runtime.cpp
index 161c4e5ec5..0d54307ef3 100644
--- a/madspace/src/python/function_runtime.cpp
+++ b/madspace/src/python/function_runtime.cpp
@@ -260,10 +260,13 @@ py::object madspace_py::tensor_to_dlpack(
     } else {
         DLDataType dtype;
         switch (tensor.dtype()) {
-        case DataType::dt_float:
+        case DataType::dt_float32:
+            dtype = {kDLFloat, 32, 1};
+            break;
+        case DataType::dt_float64:
             dtype = {kDLFloat, 64, 1};
             break;
-        case DataType::dt_int:
+        case DataType::dt_int32:
             dtype = {kDLInt, 32, 1};
             break;
         default:
@@ -410,18 +413,20 @@ Tensor madspace_py::dlpack_to_tensor(
     bool is_batch_sizes =
         expected_type ? expected_type->dtype == DataType::batch_sizes : false;
     DataType dtype;
-    if (dl_tensor->dtype.code == kDLFloat && dl_tensor->dtype.bits == 64 &&
+    if (dl_tensor->dtype.code == kDLFloat &&
+        (dl_tensor->dtype.bits == 32 || dl_tensor->dtype.bits == 64) &&
         dl_tensor->dtype.lanes == 1) {
-        dtype = DataType::dt_float;
-        if (expected_type && expected_type->dtype != DataType::dt_float) {
+        dtype =
+            dl_tensor->dtype.bits == 32 ? DataType::dt_float32 : DataType::dt_float64;
+        if (expected_type && expected_type->dtype != dtype) {
             throw std::invalid_argument(
                 std::format("Argument {}: got unexpected dtype", arg_index + 1)
             );
         }
     } else if (dl_tensor->dtype.code == kDLInt && dl_tensor->dtype.bits == 32 &&
                dl_tensor->dtype.lanes == 1) {
-        dtype = DataType::dt_int;
-        if (expected_type && expected_type->dtype != DataType::dt_int &&
+        dtype = DataType::dt_int32;
+        if (expected_type && expected_type->dtype != DataType::dt_int32 &&
             !is_batch_sizes) {
             throw std::invalid_argument(
                 std::format("Argument {}: got unexpected dtype", arg_index + 1)
@@ -439,7 +444,7 @@ Tensor madspace_py::dlpack_to_tensor(
     } else {
         throw std::invalid_argument(
             std::format(
-                "Argument {}: input dtype must be 64-bit float or 32-bit int",
+                "Argument {}: input dtype must be 32/64-bit float or 32-bit int",
                 arg_index + 1
             )
         );
@@ -532,7 +537,7 @@ Tensor madspace_py::dlpack_to_tensor(
                 std::format(
                     "Argument {}: wrong input dimension. Expected {}, got {}",
                     arg_index + 1,
-                    expected_type->shape.size() + 1,
+                    expected_type->shape.size() + is_batch,
                     dl_tensor->ndim
                 )
             );
@@ -558,7 +563,7 @@ Tensor madspace_py::dlpack_to_tensor(
         }
         if (expected_type) {
             for (int j = is_batch; j < dl_tensor->ndim; ++j) {
-                if (shape.at(j) != expected_type->shape.at(j - is_batch)) {
+                if (dl_tensor->shape[j] != expected_type->shape.at(j - is_batch)) {
                     throw std::invalid_argument(
                         std::format(
                             "Argument {}, dimension {}: shape mismatch. Expected {}, "
@@ -566,7 +571,7 @@ Tensor madspace_py::dlpack_to_tensor(
                             arg_index + 1,
                             j,
                             expected_type->shape.at(j - is_batch),
-                            shape.at(j)
+                            dl_tensor->shape[j]
                         )
                     );
                 }
diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp
index a2b615927b..fd11084261 100644
--- a/madspace/src/python/madspace.cpp
+++ b/madspace/src/python/madspace.cpp
@@ -148,8 +148,9 @@ PYBIND11_MODULE(_madspace_py, m) {
         m,
         "DataType",
         {
-            {"int", DataType::dt_int},
-            {"float", DataType::dt_float},
+            {"int32", DataType::dt_int32},
+            {"float32", DataType::dt_float32},
+            {"float64", DataType::dt_float64},
             {"batch_sizes", DataType::batch_sizes},
         },
         "",
@@ -1469,17 +1470,20 @@ PYBIND11_MODULE(_madspace_py, m) {
                std::size_t,
                std::size_t,
                MLP::Activation,
-               const std::string&>(),
+               const std::string&,
+               DataType>(),
            py::arg("input_dim"),
            py::arg("output_dim"),
            py::arg("hidden_dim") = 32,
            py::arg("layers") = 3,
            py::arg("activation") = MLP::leaky_relu,
            py::arg("prefix") = "",
+           py::arg("precision") = DataType::dt_float64,
            pydoc::doc("MLP::MLP")
     )
         .def("input_dim", &MLP::input_dim, pydoc::doc("MLP::input_dim"))
         .def("output_dim", &MLP::output_dim, pydoc::doc("MLP::output_dim"))
+        .def("precision", &MLP::precision, pydoc::doc("MLP::precision"))
         .def(
             "initialize_globals",
             &MLP::initialize_globals,
@@ -1504,7 +1508,8 @@ PYBIND11_MODULE(_madspace_py, m) {
                 std::size_t,
                 std::size_t,
                 MLP::Activation,
-                bool>(),
+                bool,
+                DataType>(),
             py::arg("input_dim"),
             py::arg("condition_dim") = 0,
             py::arg("prefix") = "",
@@ -1513,6 +1518,7 @@ PYBIND11_MODULE(_madspace_py, m) {
             py::arg("subnet_layers") = 3,
             py::arg("subnet_activation") = MLP::leaky_relu,
             py::arg("invert_spline") = true,
+            py::arg("subnet_precision") = DataType::dt_float64,
             pydoc::doc("Flow::Flow")
         )
         .def("input_dim", &Flow::input_dim, pydoc::doc("Flow::input_dim"))
@@ -1591,7 +1597,8 @@ PYBIND11_MODULE(_madspace_py, m) {
                 std::size_t,
                 MLP::Activation,
                 const std::string&,
-                bool>(),
+                bool,
+                DataType>(),
             py::arg("channel_count"),
             py::arg("particle_count"),
             py::arg("hidden_dim") = 32,
@@ -1599,6 +1606,7 @@ PYBIND11_MODULE(_madspace_py, m) {
             py::arg("activation") = MLP::leaky_relu,
             py::arg("prefix") = "",
             py::arg("include_preprocessing") = true,
+            py::arg("precision") = DataType::dt_float64,
             pydoc::doc("ChannelWeightNetwork::ChannelWeightNetwork")
         )
         .def("mlp", &ChannelWeightNetwork::mlp, pydoc::doc("ChannelWeightNetwork::mlp"))
@@ -1668,7 +1676,8 @@ PYBIND11_MODULE(_madspace_py, m) {
                 std::size_t,
                 std::size_t,
                 std::size_t,
-                MLP::Activation>(),
+                MLP::Activation,
+                DataType>(),
             py::arg("option_counts"),
             py::arg("prefix") = "",
             py::arg("dims_with_prior") = std::vector<std::size_t>{},
@@ -1676,6 +1685,7 @@ PYBIND11_MODULE(_madspace_py, m) {
             py::arg("subnet_hidden_dim") = 32,
             py::arg("subnet_layers") = 3,
             py::arg("subnet_activation") = MLP::leaky_relu,
+            py::arg("subnet_precision") = DataType::dt_float64,
             pydoc::doc("DiscreteFlow::DiscreteFlow")
         )
         .def(
diff --git a/madspace/tests/test_float32.py b/madspace/tests/test_float32.py
new file mode 100644
index 0000000000..db9022e598
--- /dev/null
+++ b/madspace/tests/test_float32.py
@@ -0,0 +1,208 @@
+"""Single-precision NN instructions: forward and backward against torch."""
+
+import pytest
+import torch
+
+import madspace as ms
+from madspace.torch import FunctionModule
+
+# a local generator, so that the global RNG state seen by other tests is unchanged
+GEN = torch.Generator().manual_seed(1234)
+
+N = 103
+DIMS_IN = 7
+DIMS_OUT = 5
+
+ACTIVATIONS = {
+    "relu": torch.relu,
+    "leaky_relu": lambda x: torch.nn.functional.leaky_relu(x, 0.01),
+    "elu": torch.nn.functional.elu,
+    "gelu": torch.nn.functional.gelu,
+    "sigmoid": torch.sigmoid,
+    "softplus": torch.nn.functional.softplus,
+    "scaled_tanh": lambda x: 4 * torch.tanh(x / 4),
+}
+
+DTYPES = {
+    "float32": (ms.DataType.float32, torch.float32, 1e-5),
+    "float64": (ms.DataType.float64, torch.float64, 1e-12),
+}
+
+
+def batch_type(dtype, shape):
+    return ms.Type(dtype, ms.batch_size, shape)
+
+
+@pytest.mark.parametrize("precision", DTYPES.keys())
+@pytest.mark.parametrize("name", ACTIVATIONS.keys())
+def test_activation(name, precision):
+    ms_dtype, torch_dtype, tol = DTYPES[precision]
+    io_type = batch_type(ms_dtype, [DIMS_IN])
+    fb = ms.FunctionBuilder(
+        ms.NamedTypes([("in", io_type)]), ms.NamedTypes([("out", io_type)])
+    )
+    fb.output(0, getattr(fb, name)(fb.input(0)))
+    module = FunctionModule(fb.function())
+
+    x = (
+        3 * torch.randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN)
+    ).requires_grad_()
+    out = module(x)
+    assert out.dtype == torch_dtype
+    out_grad = torch.randn(out.shape, dtype=out.dtype, generator=GEN)
+    (out * out_grad).sum().backward()
+
+    x_ref = x.detach().clone().requires_grad_()
+    out_ref = ACTIVATIONS[name](x_ref)
+    (out_ref * out_grad).sum().backward()
+    torch.testing.assert_close(out, out_ref, rtol=tol, atol=tol)
+    torch.testing.assert_close(x.grad, x_ref.grad, rtol=tol, atol=tol)
+
+
+@pytest.mark.parametrize("precision", DTYPES.keys())
+def test_matmul(precision):
+    ms_dtype, torch_dtype, tol = DTYPES[precision]
+    fb = ms.FunctionBuilder(
+        ms.NamedTypes(
+            [
+                ("x", batch_type(ms_dtype, [DIMS_IN])),
+                ("w", ms.Type(ms_dtype, ms.BatchSize.one, [DIMS_OUT, DIMS_IN])),
+                ("b", ms.Type(ms_dtype, ms.BatchSize.one, [DIMS_OUT])),
+            ]
+        ),
+        ms.NamedTypes([("y", batch_type(ms_dtype, [DIMS_OUT]))]),
+    )
+    fb.output(0, fb.matmul(fb.input(0), fb.input(1), fb.input(2)))
+    module = FunctionModule(fb.function())
+
+    x = torch.randn((N, DIMS_IN), dtype=torch_dtype, generator=GEN, requires_grad=True)
+    # single inputs are passed without batch dimension
+    w = torch.randn(
+        (DIMS_OUT, DIMS_IN), dtype=torch_dtype, generator=GEN, requires_grad=True
+    )
+    b = torch.randn((DIMS_OUT,), dtype=torch_dtype, generator=GEN, requires_grad=True)
+    y = module(x, w, b)
+    assert y.dtype == torch_dtype
+    y_grad = torch.randn(y.shape, dtype=y.dtype, generator=GEN)
+    (y * y_grad).sum().backward()
+
+    refs = [t.detach().clone().requires_grad_() for t in (x, w, b)]
+    y_ref = refs[0] @ refs[1].T + refs[2]
+    (y_ref * y_grad).sum().backward()
+    torch.testing.assert_close(y, y_ref, rtol=tol, atol=tol)
+    for t, t_ref in zip((x, w, b), refs):
+        torch.testing.assert_close(t.grad, t_ref.grad, rtol=tol, atol=tol)
+
+
+def test_convert_roundtrip():
+    in_type = batch_type(ms.DataType.float64, [3])
+    fb = ms.FunctionBuilder(
+        ms.NamedTypes([("in", in_type)]),
+        ms.NamedTypes(
+            [("single", batch_type(ms.DataType.float32, [3])), ("out", in_type)]
+        ),
+    )
+    single = fb.convert_float32(fb.input(0))
+    fb.output(0, single)
+    fb.output(1, fb.convert_float64(fb.relu(single)))
+    module = FunctionModule(fb.function())
+
+    x = torch.randn((N, 3), dtype=torch.float64, generator=GEN, requires_grad=True)
+    single_out, out = module(x)
+    assert single_out.dtype == torch.float32
+    assert out.dtype == torch.float64
+    torch.testing.assert_close(single_out, x.detach().float())
+    torch.testing.assert_close(out, torch.relu(x.detach().float()).double())
+    out.sum().backward()
+    torch.testing.assert_close(x.grad, (x.detach() > 0).double())
+
+
+def test_literal_converted_for_float32():
+    io_type = batch_type(ms.DataType.float32, [1])
+    fb = ms.FunctionBuilder(
+        ms.NamedTypes([("in", io_type)]), ms.NamedTypes([("out", io_type)])
+    )
+    # float64 literals are converted when the instruction dispatches to float32
+    fb.output(0, fb.matmul(fb.input(0), ms.Value([2.0], [1, 1]), ms.Value([1.0])))
+    module = FunctionModule(fb.function())
+    x = torch.randn((N, 1), dtype=torch.float32, generator=GEN)
+    torch.testing.assert_close(module(x), 2 * x + 1)
+
+
+def test_flow_float32_subnets(tmp_path):
+    """A flow with float32 subnetworks matches the float64 one with the same
+    (rounded) weights; the Jacobian and the output stay in float64."""
+    # a private context, so that no globals are left in the default one
+    ctx = ms.Context()
+    flows = {
+        dtype: ms.Flow(
+            4, prefix=f"fp32_test_{name}", subnet_precision=dtype, invert_spline=False
+        )
+        for name, dtype in [("f32", ms.DataType.float32), ("f64", ms.DataType.float64)]
+    }
+    for flow in flows.values():
+        flow.initialize_globals(ctx, 42)
+    # randomize all weights (the last layers start at zero) and round to float32
+    names_f64 = [n for n in ctx.global_names() if n.startswith("fp32_test_f64")]
+    for name in names_f64:
+        weights = ctx.get_global(name).torch()
+        weights[:] = (
+            0.3
+            * torch.randn(weights.shape, generator=GEN, dtype=torch.float32).double()
+        )
+        ctx.get_global(name.replace("f64", "f32")).torch()[:] = weights.float()
+
+    # outputs and gradients with respect to the input and the float32 weights
+    r = torch.rand((N, 4), dtype=torch.float64, generator=GEN)
+    outputs, grads = {}, {}
+    for dtype, flow in flows.items():
+        module = FunctionModule(flow.forward_function(), ctx)
+        r_grad = r.clone().requires_grad_()
+        out, det = module(r_grad)
+        assert out.dtype == det.dtype == torch.float64
+        (out.sum() + det.log().sum()).backward()
+        weight = next(p for n, p in module.global_params.items() if "weight" in n)
+        assert weight.grad.dtype == weight.dtype
+        outputs[dtype] = (out.detach(), det.detach())
+        grads[dtype] = (r_grad.grad, weight.grad.double())
+    for o32, o64 in zip(outputs[ms.DataType.float32], outputs[ms.DataType.float64]):
+        torch.testing.assert_close(o32, o64, rtol=1e-4, atol=1e-5)
+    for g32, g64 in zip(grads[ms.DataType.float32], grads[ms.DataType.float64]):
+        torch.testing.assert_close(g32, g64, rtol=1e-3, atol=1e-4)
+
+    # float32 globals survive saving and loading
+    ctx.save_globals(str(tmp_path / "globals.tar"))
+    loaded = ms.Context()
+    loaded.load_globals(str(tmp_path / "globals.tar"))
+    for name in names_f64:
+        name32 = name.replace("f64", "f32")
+        global32 = loaded.get_global(name32).torch()
+        assert global32.dtype == torch.float32
+        torch.testing.assert_close(global32, ctx.get_global(name32).torch())
+
+
+def test_structural_ops_backward():
+    """stack, unstack, cat, batch_cat and batch_split propagate float32 gradients."""
+    in_type = batch_type(ms.DataType.float32, [3])
+    fb = ms.FunctionBuilder(
+        ms.NamedTypes([("in", in_type)]),
+        ms.NamedTypes([("out", batch_type(ms.DataType.float32, [12]))]),
+    )
+    x = fb.input(0)
+    a, b = fb.unstack(fb.stack([x, fb.relu(x)]))
+    c = fb.cat([a, b])
+    combined, counts = fb.batch_cat([c, c])
+    fb.output(0, fb.cat(fb.batch_split(combined, counts)))
+    module = FunctionModule(fb.function())
+
+    x = torch.randn((N, 3), dtype=torch.float32, generator=GEN, requires_grad=True)
+    out = module(x)
+    out_grad = torch.randn(out.shape, dtype=out.dtype, generator=GEN)
+    (out * out_grad).sum().backward()
+
+    x_ref = x.detach().clone().requires_grad_()
+    c_ref = torch.cat([x_ref, torch.relu(x_ref)], dim=1)
+    out_ref = torch.cat([c_ref, c_ref], dim=1)
+    (out_ref * out_grad).sum().backward()
+    torch.testing.assert_close(out, out_ref)
+    torch.testing.assert_close(x.grad, x_ref.grad)
diff --git a/madspace/tests/test_histogram_op.py b/madspace/tests/test_histogram_op.py
index f186cd91f4..77f5073838 100644
--- a/madspace/tests/test_histogram_op.py
+++ b/madspace/tests/test_histogram_op.py
@@ -16,7 +16,7 @@
 
 
 def histogram_function(n_bins):
-    out_type = ms.Type(ms.DataType.float, ms.BatchSize.one, [n_bins + 2])
+    out_type = ms.Type(ms.DataType.float64, ms.BatchSize.one, [n_bins + 2])
     fb = ms.FunctionBuilder(
         ms.NamedTypes([(name, ms.batch_float) for name in ("x", "w", "lo", "hi")]),
         ms.NamedTypes([("values", out_type), ("square_values", out_type)]),
diff --git a/madspace/tests/test_mlp.py b/madspace/tests/test_mlp.py
index 8f46de582c..c496c05a95 100644
--- a/madspace/tests/test_mlp.py
+++ b/madspace/tests/test_mlp.py
@@ -14,9 +14,20 @@ def assert_approx(actual, expected):
     torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-12)
 
 
+PRECISIONS = {
+    "float64": (ms.DataType.float64, torch.float64, 1e-6),
+    "float32": (ms.DataType.float32, torch.float32, 1e-3),
+}
+
+
+@pytest.fixture(params=PRECISIONS.keys())
+def precision(request):
+    return PRECISIONS[request.param]
+
+
 @pytest.fixture
-def mlp():
-    return ms.MLP(10, 1, 32, 3, ms.MLP.leaky_relu, "")
+def mlp(precision):
+    return ms.MLP(10, 1, 32, 3, ms.MLP.leaky_relu, "", precision[0])
 
 
 @pytest.fixture
@@ -34,9 +45,16 @@ def test_properties():
     mlp = ms.MLP(10, 1, 32, 3, ms.MLP.leaky_relu, "")
     assert mlp.input_dim() == 10
     assert mlp.output_dim() == 1
+    assert mlp.precision() == ms.DataType.float64
+    mlp = ms.MLP(10, 1, precision=ms.DataType.float32)
+    assert mlp.precision() == ms.DataType.float32
+    with pytest.raises(Exception):
+        ms.MLP(10, 1, precision=ms.DataType.int32)
 
 
-def test_initialization(mlp_module):
+def test_initialization(mlp_module, precision):
+    for param in mlp_module.global_params.values():
+        assert param.dtype == precision[1]
     assert torch.all(mlp_module.global_params["layer1:weight"] != 0)
     assert torch.all(mlp_module.global_params["layer1:bias"] != 0)
     assert torch.all(mlp_module.global_params["layer2:weight"] != 0)
@@ -73,14 +91,16 @@ def test_activation(activation):
     assert_approx(grad_me, grad_torch)
 
 
-def test_training(mlp, mlp_module, module_context):
+@pytest.mark.parametrize("grad_clip", [0.0, 0.01])
+def test_training(mlp, mlp_module, module_context, precision, grad_clip):
+    _, dtype, tol = precision
     mlp_torch = nn.Sequential(
         nn.Linear(10, 32),
         nn.LeakyReLU(),
         nn.Linear(32, 32),
         nn.LeakyReLU(),
         nn.Linear(32, 1),
-    )
+    ).to(dtype)
 
     with torch.no_grad():
         for i in range(3):
@@ -102,47 +122,61 @@ def test_training(mlp, mlp_module, module_context):
     fb.output(1, mlp_out)
     ctx.copy_globals_from(module_context)
     builtin_opt = ms.AdamOptimizer(
-        function=fb.function(), context=ctx, learning_rate=1e-2
+        function=fb.function(),
+        context=ctx,
+        learning_rate=1e-2,
+        grad_clip_threshold=grad_clip,
     )
 
     opt_ms = torch.optim.Adam(mlp_module.parameters(), lr=1e-2)
     opt_torch = torch.optim.Adam(mlp_torch.parameters(), lr=1e-2)
 
+    def close(a, b):
+        return a == approx(b, rel=tol, abs=tol)
+
+    def assert_close(actual, expected):
+        torch.testing.assert_close(actual, expected, rtol=tol, atol=tol)
+
+    # the builtin optimizer clips the RMS gradient, not the norm
+    n_params = sum(p.numel() for p in mlp_torch.parameters())
+    max_norm = grad_clip * n_params**0.5
+
     g = generator(3210)
     for i in range(10):
-        x = randn((128, 10), generator=g)
+        # the input and the loss are in float64, only the network uses dtype
+        x = randn((128, 10), generator=g, dtype=torch.float64)
         y = x.square().sum(dim=1, keepdim=True)
 
         loss_ms = (y - mlp_module(x)).square().mean()
         opt_ms.zero_grad()
         loss_ms.backward()
+        if grad_clip > 0:
+            clipped = nn.utils.clip_grad_norm_(mlp_module.parameters(), max_norm)
+            assert clipped > max_norm
         opt_ms.step()
 
-        loss_torch = (y - mlp_torch(x)).square().mean()
+        loss_torch = (y - mlp_torch(x.to(dtype)).double()).square().mean()
         opt_torch.zero_grad()
         loss_torch.backward()
+        if grad_clip > 0:
+            nn.utils.clip_grad_norm_(mlp_torch.parameters(), max_norm)
         opt_torch.step()
 
         loss_builtin, mlp_out = builtin_opt.step([x])
 
-        assert loss_ms.item() == approx(loss_torch.item())
-        assert loss_builtin.torch().item() == approx(loss_torch.item())
+        assert close(loss_ms.item(), loss_torch.item())
+        assert close(loss_builtin.torch().item(), loss_torch.item())
 
         with torch.no_grad():
             for i in range(3):
                 w_torch = mlp_torch[2 * i].weight
                 b_torch = mlp_torch[2 * i].bias
-                assert_approx(
-                    mlp_module.global_params[f"layer{i+1}:weight"][0], w_torch
-                )
-                assert_approx(mlp_module.global_params[f"layer{i+1}:bias"][0], b_torch)
-                assert_approx(
-                    mlp_module.global_params[f"layer{i+1}:weight"].grad[0],
-                    mlp_torch[2 * i].weight.grad,
-                )
-                assert_approx(
-                    mlp_module.global_params[f"layer{i+1}:bias"].grad[0],
-                    mlp_torch[2 * i].bias.grad,
-                )
-                assert_approx(ctx.get_global(f"layer{i+1}.weight").torch()[0], w_torch)
-                assert_approx(ctx.get_global(f"layer{i+1}.bias").torch()[0], b_torch)
+                w_ms = mlp_module.global_params[f"layer{i+1}:weight"]
+                b_ms = mlp_module.global_params[f"layer{i+1}:bias"]
+                assert w_ms.dtype == dtype
+                assert_close(w_ms[0], w_torch)
+                assert_close(b_ms[0], b_torch)
+                assert_close(w_ms.grad[0], mlp_torch[2 * i].weight.grad)
+                assert_close(b_ms.grad[0], mlp_torch[2 * i].bias.grad)
+                assert_close(ctx.get_global(f"layer{i+1}.weight").torch()[0], w_torch)
+                assert_close(ctx.get_global(f"layer{i+1}.bias").torch()[0], b_torch)
diff --git a/madspace/tests/test_rank4_kernels.py b/madspace/tests/test_rank4_kernels.py
index e31c6a2ed5..37b9bb46cb 100644
--- a/madspace/tests/test_rank4_kernels.py
+++ b/madspace/tests/test_rank4_kernels.py
@@ -11,7 +11,7 @@
 
 
 def test_rank4_add(n=4096):
-    rank4 = ms.Type(ms.DataType.float, ms.batch_size, [2, 3, 4])
+    rank4 = ms.Type(ms.DataType.float64, ms.batch_size, [2, 3, 4])
     fb = ms.FunctionBuilder(
         ms.NamedTypes([("x", rank4)]), ms.NamedTypes([("y", rank4)])
     )
diff --git a/madspace/tests/test_simd_layout.py b/madspace/tests/test_simd_layout.py
index 37f592f576..8ef4a03494 100644
--- a/madspace/tests/test_simd_layout.py
+++ b/madspace/tests/test_simd_layout.py
@@ -97,7 +97,7 @@ def body(fb, x):
     func = build(
         [("x", ms.batch_float_array(3))], [("y", ms.batch_float_array(3))], body
     )
-    assert instruction_names(func) == ["mul"]
+    assert instruction_names(func) == ["mul_float64"]
 
 
 @pytest.mark.cpu_only
@@ -150,7 +150,7 @@ def test_single_flavor_integrand_has_no_constant_gathers(tmp_path):
             for output in instr.outputs:
                 sources[output.local_index] = instr.instruction.name
         for instr in func.instructions:
-            if instr.instruction.name == "mul":
+            if instr.instruction.name == "mul_float64":
                 assert all(
                     sources.get(arg.local_index) != "full" for arg in instr.inputs
                 )
@@ -276,7 +276,7 @@ def test_global_probability_gradient_is_summed_over_events():
     """The gradient of a broadcast (global) probability vector is the sum of the
     contributions of all events."""
     ctx = new_context()
-    ctx.define_global("test_simd_layout_probs", ms.DataType.float, [4], True)
+    ctx.define_global("test_simd_layout_probs", ms.DataType.float64, [4], True)
     ctx.get_global("test_simd_layout_probs").torch()[:] = torch.tensor(
         [[0.1, 0.4, 0.3, 0.2]]
     )
@@ -284,7 +284,7 @@ def test_global_probability_gradient_is_summed_over_events():
         ms.NamedTypes([("index", ms.batch_int)]),
         ms.NamedTypes([("r", ms.batch_float), ("det", ms.batch_float)]),
     )
-    probs = fb.get_global("test_simd_layout_probs", ms.DataType.float, [4])
+    probs = fb.get_global("test_simd_layout_probs", ms.DataType.float64, [4])
     r, det = fb.sample_discrete_probs_inverse(fb.input(0), probs)
     fb.output(0, r)
     fb.output(1, det)
@@ -317,13 +317,13 @@ def test_broadcast_gradient_through_stack():
     """stack of a broadcast (global) and a batched value: the gradient of the
     global accumulates the whole batch."""
     ctx = new_context()
-    ctx.define_global("test_simd_layout_scalar", ms.DataType.float, [], True)
+    ctx.define_global("test_simd_layout_scalar", ms.DataType.float64, [], True)
     ctx.get_global("test_simd_layout_scalar").torch()[:] = 1.5
     fb = ms.FunctionBuilder(
         ms.NamedTypes([("x", ms.batch_float)]),
         ms.NamedTypes([("y", ms.batch_float_array(2))]),
     )
-    g = fb.get_global("test_simd_layout_scalar", ms.DataType.float, [])
+    g = fb.get_global("test_simd_layout_scalar", ms.DataType.float64, [])
     fb.output(0, fb.stack([g, fb.input(0)]))
     module = FunctionModule(fb.function(), ctx)
     param = module.global_params["test_simd_layout_scalar"]
diff --git a/madspace/tests/test_tar.py b/madspace/tests/test_tar.py
index d16d8ae093..8fd8ec1ede 100644
--- a/madspace/tests/test_tar.py
+++ b/madspace/tests/test_tar.py
@@ -19,7 +19,7 @@ def make_context(tmp_path):
         "d.scalar": rng.normal(size=(1, 1)),
     }
     for name, value in values.items():
-        dtype = ms.DataType.int if value.dtype == np.int32 else ms.DataType.float
+        dtype = ms.DataType.int32 if value.dtype == np.int32 else ms.DataType.float64
         glob = ctx.define_global(name, dtype, list(value.shape[1:]))
         glob.torch().copy_(torch.from_numpy(value))
     return ctx, values
diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml
index cbf2cc6bcf..b550825e6a 100644
--- a/tests/input_files/mg7_run_card.toml
+++ b/tests/input_files/mg7_run_card.toml
@@ -109,6 +109,9 @@ cwnet_hidden_dim = 64
 cwnet_layers = 3
 cwnet_activation = "leaky_relu" # options: relu, leaky_relu, elu, gelu, sigmoid, softplus
 
+# precision of all networks
+nn_float_precision = "float64" # options: float32, float64
+
 # training parameters
 train_batches = 1000
 log_interval = 100
diff --git a/tests/unit_tests/various/test_banner.py b/tests/unit_tests/various/test_banner.py
index b7d92964ad..4020d65055 100755
--- a/tests/unit_tests/various/test_banner.py
+++ b/tests/unit_tests/various/test_banner.py
@@ -1985,6 +1985,7 @@ def test_defaults_and_section_access(self):
         self.assertEqual(rc['beam']['e_cm'], 13000.0)
         self.assertIs(rc['vegas']['enable'], True)
         self.assertEqual(rc['madnis']['enable'], 'auto')
+        self.assertEqual(rc['madnis']['nn_float_precision'], 'float64')
         self.assertEqual(rc['phasespace']['mode'], 'multichannel')
         # the "enable" key colliding between [vegas] and [madnis] is stored
         # independently, so editing one does not affect the other
@@ -2013,17 +2014,20 @@ def test_write_is_valid_toml(self):
         self.assertEqual(data['multiparticles']['photon'], [22])
         self.assertEqual(data['cuts']['sqrt_s'], {'min': 0.0})
         self.assertEqual(data['madnis']['adam_eps'], 1e-8)
+        self.assertEqual(data['madnis']['nn_float_precision'], 'float64')
 
     def test_round_trip(self):
         """write() then read() preserves every value"""
         rc = bannermod.RunCardMG7()
         rc['generation']['events'] = 250
         rc['phasespace']['mode'] = 'flat'
+        rc['madnis']['nn_float_precision'] = 'float32'
         out = io.StringIO()
         rc.write(out, template=self.template)
         rc2 = bannermod.RunCardMG7(out.getvalue())
         self.assertEqual(rc2['generation']['events'], 250)
         self.assertEqual(rc2['phasespace']['mode'], 'flat')
+        self.assertEqual(rc2['madnis']['nn_float_precision'], 'float32')
         self.assertEqual(rc2['multiparticles']['lepton'], rc['multiparticles']['lepton'])
         self.assertEqual(rc2['cuts']['jet-delta_r'], {'min': 0.4})