diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst
index b310007ac8..8b44707844 100644
--- a/docs/source/run_card.rst
+++ b/docs/source/run_card.rst
@@ -616,3 +616,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`` (``"float32"``)
+    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 6c70fbdbeb..015ddc8df9 100644
--- a/madgraph/iolibs/template_files/mg7/launch.py
+++ b/madgraph/iolibs/template_files/mg7/launch.py
@@ -3507,6 +3507,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)
@@ -3529,6 +3530,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
@@ -3543,6 +3545,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)
 
@@ -3623,6 +3626,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
@@ -3640,6 +3644,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 42bcc709af..60c930045f 100644
--- a/madgraph/iolibs/template_files/mg7/run_card.toml
+++ b/madgraph/iolibs/template_files/mg7/run_card.toml
@@ -191,6 +191,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 eaa49f8703..c0c1de446b 100755
--- a/madgraph/various/banner.py
+++ b/madgraph/various/banner.py
@@ -7192,6 +7192,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', "float32",
+            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 c7989a423f..f4ff1f2952 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..26560101d8 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(
@@ -249,28 +311,45 @@ def instruction_set_mixin(commands, out_dir):
         f_op.write("\n")
 
 
+def custom_op_func(prefix, cmd):
+    """The C++ function of a custom op. custom_op can name an implementation shared
+    by several instructions instead of True, with the leading template arguments
+    custom_op_args."""
+    custom_op = cmd["custom_op"]
+    name = custom_op if isinstance(custom_op, str) else cmd["base_name"]
+    template_args = list(cmd.get("custom_op_args", []))
+    if "ctype" in cmd:
+        template_args.append(cmd["ctype"])
+    func = f"{prefix}{name}"
+    return f"{func}<{', '.join(template_args)}>" if template_args else func
+
+
 def runtime_mixin(commands, device, out_dir):
     with open(out_dir / device / "runtime_mixin.inc", "w") as f:
         write_autogen(f)
 
         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 = custom_op_func("op_", cmd)
             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 +372,13 @@ 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 = custom_op_func("backward_op_", cmd)
                 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 +400,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 +432,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 +449,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 +460,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 b83c686e40..afa9547d43 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",
             "--enum-class-locations=JetScaleScheme:MLMClustering.JetScaleScheme",
             "--enum-class-locations=ScaleScheme:MLMClustering.ScaleScheme",
             "--enum-class-locations=PartonLineScheme:MLMClustering.PartonLineScheme",
diff --git a/madspace/include/madspace/compgraphs/function.hpp b/madspace/include/madspace/compgraphs/function.hpp
index 42c0a48dd2..25ba67fede 100644
--- a/madspace/include/madspace/compgraphs/function.hpp
+++ b/madspace/include/madspace/compgraphs/function.hpp
@@ -81,6 +81,7 @@ class Function {
     std::vector<InstructionCall> _instructions;
 
     friend Function sort_breadth_first(const Function& function);
+    friend Function fuse_instructions(const Function& function);
 };
 
 std::ostream& operator<<(std::ostream& out, const Value& value);
@@ -128,6 +129,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/optimizer.hpp b/madspace/include/madspace/compgraphs/optimizer.hpp
index a0dd35a775..546902f59a 100644
--- a/madspace/include/madspace/compgraphs/optimizer.hpp
+++ b/madspace/include/madspace/compgraphs/optimizer.hpp
@@ -33,4 +33,20 @@ class LastUseOfLocals {
 
 Function sort_breadth_first(const Function& function);
 
+/**
+ * Lowering pass for networks, run by the GPU backend when it builds a runtime.
+ * Fuses an element-wise instruction into the instruction that produces its input:
+ *  - `convert_float32(v)` where `v` is an output of `select`, `rqs_forward` or
+ *    `rqs_inverse` becomes the `*_with_float32` variant of the producer, which
+ *    writes the single-precision copy itself;
+ *  - `convert_float64(matmul_float32(x, w, b))`, with the matmul result used only
+ *    by the conversion, becomes `matmul_to_float64`;
+ *  - an activation `act(matmul(x, w, b))` (relu, leaky_relu, elu, gelu, sigmoid or
+ *    softplus, in either precision) becomes `matmul_act`, which also keeps the
+ *    matmul result as an output.
+ * The graph keeps only the generic instructions and so stays portable; the fused
+ * instructions never appear in a saved function.
+ */
+Function fuse_instructions(const Function& function);
+
 } // namespace madspace
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 e9368d607c..bcb1e3f807 100644
--- a/madspace/include/madspace/driver/tensor.hpp
+++ b/madspace/include/madspace/driver/tensor.hpp
@@ -448,11 +448,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));
@@ -463,11 +459,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));
@@ -481,8 +473,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)
             ),
@@ -532,6 +524,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());
     }
@@ -541,6 +534,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());
     }
@@ -549,6 +543,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());
@@ -634,9 +629,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;
@@ -705,6 +702,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.
@@ -809,7 +809,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.
@@ -923,6 +923,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 f78a3c0ee1..c099e3cbe4 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, ...]
@@ -3162,7 +3163,53 @@ 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_to_float64:
+  inputs:
+    - name: x
+      type: [float32, n]
+      backward_arg: True
+    - name: weight
+      type: [float32, single, m, n]
+      backward_arg: True
+    - name: bias
+      type: [float32, single, m]
+  outputs:
+    - name: y
+      type: [float64, m]
+  custom_op: True
+  differentiable: True
+  desc: "Single-precision affine map whose result is converted to double precision. Equivalent to @ref convert_float64 applied to the output of @ref matmul, but the GPU backend fuses the bias addition and the conversion, and the backward pass fuses the gradient conversion with the bias gradient. Used for the last layer of single-precision @ref MLP networks."
+
 matmul:
+  precision: [float32, float64]
   inputs:
     - name: x
       type: [float, n]
@@ -3183,7 +3230,164 @@ matmul:
   differentiable: True
   desc: "Affine layer y = weight x + bias. Building block of @ref MLP."
 
+matmul_relu:
+  precision: [float32, float64]
+  inputs:
+    - name: x
+      type: [float, n]
+      desc: "input vector"
+      backward_arg: True
+    - name: weight
+      type: [float, single, m, n]
+      desc: "weight matrix"
+      backward_arg: True
+    - name: bias
+      type: [float, single, m]
+      desc: "bias vector"
+  outputs:
+    - name: linear
+      type: [float, m]
+      desc: "the affine transform of the input"
+    - name: y
+      type: [float, m]
+      desc: "the activated affine transform"
+  custom_op: matmul_activation
+  custom_op_args: [ReluActivation]
+  differentiable: True
+  desc: "@ref matmul followed by @ref relu, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref relu applied to the output of @ref matmul."
+
+matmul_leaky_relu:
+  precision: [float32, float64]
+  inputs:
+    - name: x
+      type: [float, n]
+      desc: "input vector"
+      backward_arg: True
+    - name: weight
+      type: [float, single, m, n]
+      desc: "weight matrix"
+      backward_arg: True
+    - name: bias
+      type: [float, single, m]
+      desc: "bias vector"
+  outputs:
+    - name: linear
+      type: [float, m]
+      desc: "the affine transform of the input"
+    - name: y
+      type: [float, m]
+      desc: "the activated affine transform"
+  custom_op: matmul_activation
+  custom_op_args: [LeakyReluActivation]
+  differentiable: True
+  desc: "@ref matmul followed by @ref leaky_relu, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref leaky_relu applied to the output of @ref matmul."
+
+matmul_elu:
+  precision: [float32, float64]
+  inputs:
+    - name: x
+      type: [float, n]
+      desc: "input vector"
+      backward_arg: True
+    - name: weight
+      type: [float, single, m, n]
+      desc: "weight matrix"
+      backward_arg: True
+    - name: bias
+      type: [float, single, m]
+      desc: "bias vector"
+  outputs:
+    - name: linear
+      type: [float, m]
+      desc: "the affine transform of the input"
+    - name: y
+      type: [float, m]
+      desc: "the activated affine transform"
+  custom_op: matmul_activation
+  custom_op_args: [EluActivation]
+  differentiable: True
+  desc: "@ref matmul followed by @ref elu, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref elu applied to the output of @ref matmul."
+
+matmul_gelu:
+  precision: [float32, float64]
+  inputs:
+    - name: x
+      type: [float, n]
+      desc: "input vector"
+      backward_arg: True
+    - name: weight
+      type: [float, single, m, n]
+      desc: "weight matrix"
+      backward_arg: True
+    - name: bias
+      type: [float, single, m]
+      desc: "bias vector"
+  outputs:
+    - name: linear
+      type: [float, m]
+      desc: "the affine transform of the input"
+    - name: y
+      type: [float, m]
+      desc: "the activated affine transform"
+  custom_op: matmul_activation
+  custom_op_args: [GeluActivation]
+  differentiable: True
+  desc: "@ref matmul followed by @ref gelu, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref gelu applied to the output of @ref matmul."
+
+matmul_sigmoid:
+  precision: [float32, float64]
+  inputs:
+    - name: x
+      type: [float, n]
+      desc: "input vector"
+      backward_arg: True
+    - name: weight
+      type: [float, single, m, n]
+      desc: "weight matrix"
+      backward_arg: True
+    - name: bias
+      type: [float, single, m]
+      desc: "bias vector"
+  outputs:
+    - name: linear
+      type: [float, m]
+      desc: "the affine transform of the input"
+    - name: y
+      type: [float, m]
+      desc: "the activated affine transform"
+  custom_op: matmul_activation
+  custom_op_args: [SigmoidActivation]
+  differentiable: True
+  desc: "@ref matmul followed by @ref sigmoid, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref sigmoid applied to the output of @ref matmul."
+
+matmul_softplus:
+  precision: [float32, float64]
+  inputs:
+    - name: x
+      type: [float, n]
+      desc: "input vector"
+      backward_arg: True
+    - name: weight
+      type: [float, single, m, n]
+      desc: "weight matrix"
+      backward_arg: True
+    - name: bias
+      type: [float, single, m]
+      desc: "bias vector"
+  outputs:
+    - name: linear
+      type: [float, m]
+      desc: "the affine transform of the input"
+    - name: y
+      type: [float, m]
+      desc: "the activated affine transform"
+  custom_op: matmul_activation
+  custom_op_args: [SoftplusActivation]
+  differentiable: True
+  desc: "@ref matmul followed by @ref softplus, with the bias addition and the activation in one kernel. Produced by the GPU backend's lowering of @ref softplus applied to the output of @ref matmul."
+
 relu:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -3197,6 +3401,7 @@ relu:
   desc: "Element-wise ReLU activation. Building block of @ref MLP."
 
 leaky_relu:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -3210,6 +3415,7 @@ leaky_relu:
   desc: "Element-wise leaky ReLU activation. Building block of @ref MLP."
 
 elu:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -3223,6 +3429,7 @@ elu:
   desc: "Element-wise ELU activation. Building block of @ref MLP."
 
 gelu:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -3236,6 +3443,7 @@ gelu:
   desc: "Element-wise GELU activation. Building block of @ref MLP."
 
 sigmoid:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -3250,6 +3458,7 @@ sigmoid:
   desc: "Element-wise sigmoid activation. Building block of @ref MLP."
 
 scaled_tanh:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -3264,6 +3473,7 @@ scaled_tanh:
   desc: "Element-wise scaled hyperbolic tangent activation. Building block of @ref MLP."
 
 softplus:
+  precision: [float32, float64]
   inputs:
     - name: in1
       type: [float, ...]
@@ -3354,6 +3564,52 @@ rqs_inverse:
   dims: 2
   desc: "Applies a rational-quadratic spline in the inverse direction. Building block of @ref Flow."
 
+rqs_forward_with_float32:
+  inputs:
+    - name: input
+      type: [float, n]
+      desc: "values to transform"
+    - name: condition
+      type: [float, n, 6]
+      desc: "packed spline parameters"
+  outputs:
+    - name: output
+      type: [float, n]
+      desc: "transformed values"
+    - name: det
+      type: [float, n]
+      desc: "per-element Jacobian"
+    - name: output32
+      type: [float32, n]
+      desc: "the transformed values in single precision"
+  differentiable: True
+  dims: 2
+  vectorized: False
+  desc: "@ref rqs_forward that also returns its output in single precision, so that the conversion does not need a separate kernel. Produced by the GPU backend's lowering of @ref convert_float32 applied to the output of @ref rqs_forward."
+
+rqs_inverse_with_float32:
+  inputs:
+    - name: input
+      type: [float, n]
+      desc: "values to transform"
+    - name: condition
+      type: [float, n, 6]
+      desc: "packed spline parameters"
+  outputs:
+    - name: output
+      type: [float, n]
+      desc: "transformed values"
+    - name: det
+      type: [float, n]
+      desc: "per-element Jacobian"
+    - name: output32
+      type: [float32, n]
+      desc: "the transformed values in single precision"
+  differentiable: True
+  dims: 2
+  vectorized: False
+  desc: "@ref rqs_inverse that also returns its output in single precision, so that the conversion does not need a separate kernel. Produced by the GPU backend's lowering of @ref convert_float32 applied to the output of @ref rqs_inverse."
+
 softmax:
   inputs:
     - name: input
@@ -3570,6 +3826,26 @@ select:
   differentiable: True
   desc: "Gathers several entries from a float vector by index."
 
+select_with_float32:
+  inputs:
+    - name: input
+      type: [float, n]
+      desc: "source vector"
+    - name: indices
+      type: [int, single, m]
+      desc: "indices to gather"
+      backward_arg: True
+  outputs:
+    - name: output
+      type: [float, m]
+      desc: "gathered entries"
+    - name: output32
+      type: [float32, m]
+      desc: "the gathered entries in single precision"
+  differentiable: True
+  vectorized: False
+  desc: "@ref select that also returns the gathered entries in single precision, so that the conversion does not need a separate kernel. Produced by the GPU backend's lowering of @ref convert_float32 applied to the output of @ref select."
+
 select_vector:
   inputs:
     - name: input
diff --git a/madspace/madspace/python_function.py b/madspace/madspace/python_function.py
index d2b11b47bd..7c4ae113d4 100644
--- a/madspace/madspace/python_function.py
+++ b/madspace/madspace/python_function.py
@@ -12,33 +12,33 @@
 
 # input name: (data type, shape after the batch dimension); "n" is the particle count
 _INPUT_SHAPES = {
-    "momenta": (DataType.float, ("n", 4)),
-    "alpha_s": (DataType.float, ()),
-    "flavor_index": (DataType.int, ()),
-    "random_color": (DataType.float, ()),
-    "random_helicity": (DataType.float, ()),
-    "random_diagram": (DataType.float, ()),
-    "helicity_index": (DataType.int, ()),
-    "diagram_index": (DataType.int, ()),
-    "channel_index": (DataType.int, ()),
-    "particle_count": (DataType.int, ()),
-    "pdg_ids": (DataType.int, ("n",)),
-    "momentum_fraction": (DataType.float, ()),
-    "scale": (DataType.float, ()),
-    "pdg_id": (DataType.int, ()),
-    "beam_index": (DataType.int, ()),
+    "momenta": (DataType.float64, ("n", 4)),
+    "alpha_s": (DataType.float64, ()),
+    "flavor_index": (DataType.int32, ()),
+    "random_color": (DataType.float64, ()),
+    "random_helicity": (DataType.float64, ()),
+    "random_diagram": (DataType.float64, ()),
+    "helicity_index": (DataType.int32, ()),
+    "diagram_index": (DataType.int32, ()),
+    "channel_index": (DataType.int32, ()),
+    "particle_count": (DataType.int32, ()),
+    "pdg_ids": (DataType.int32, ("n",)),
+    "momentum_fraction": (DataType.float64, ()),
+    "scale": (DataType.float64, ()),
+    "pdg_id": (DataType.int32, ()),
+    "beam_index": (DataType.int32, ()),
 }
 _OUTPUT_TYPES = {
-    "matrix_element": DataType.float,
-    "color_index": DataType.int,
-    "helicity_index": DataType.int,
-    "diagram_index": DataType.int,
-    "cut": DataType.float,
-    "ren_scale": DataType.float,
-    "fact_scale1": DataType.float,
-    "fact_scale2": DataType.float,
-    "pdf": DataType.float,
-    "bias": DataType.float,
+    "matrix_element": DataType.float64,
+    "color_index": DataType.int32,
+    "helicity_index": DataType.int32,
+    "diagram_index": DataType.int32,
+    "cut": DataType.float64,
+    "ren_scale": DataType.float64,
+    "fact_scale1": DataType.float64,
+    "fact_scale2": DataType.float64,
+    "pdf": DataType.float64,
+    "bias": DataType.float64,
 }
 
 
@@ -135,7 +135,7 @@ def _view(self, address, dtype, shape, stride, offset, device):
 
     def _numpy_output(self, address, dtype, count, offset):
         # numpy.from_dlpack gives read-only arrays, so map the CPU buffer directly
-        ctype = ctypes.c_double if dtype == DataType.float else ctypes.c_int32
+        ctype = ctypes.c_double if dtype == DataType.float64 else ctypes.c_int32
         size = ctypes.sizeof(ctype)
         buffer = (ctype * count).from_address(address + offset * size)
         return self._numpy.frombuffer(buffer, dtype=self._numpy.dtype(ctype))
@@ -169,7 +169,7 @@ def _call(
             if self._needs_particle_count and n is None:
                 counts = self._view(
                     inputs[_UMAMI_INPUT_KEYS["particle_count"]],
-                    DataType.int,
+                    DataType.int32,
                     [count],
                     [1],
                     offset,
@@ -227,7 +227,7 @@ def _check(self, key, result, dtype, count, device):
                 raise TypeError(
                     f"{prefix} must be a torch.Tensor, not {type(result).__name__}"
                 )
-            expected = torch.float64 if dtype == DataType.float else torch.int32
+            expected = torch.float64 if dtype == DataType.float64 else torch.int32
             result_device = "cpu" if result.device.type == "cpu" else "gpu"
             if result_device != ("cpu" if device[0] == "cpu" else "gpu"):
                 raise ValueError(
@@ -240,7 +240,7 @@ def _check(self, key, result, dtype, count, device):
                     f"{prefix} must be a numpy.ndarray, not {type(result).__name__}"
                 )
             expected = numpy.dtype(
-                numpy.float64 if dtype == DataType.float else numpy.int32
+                numpy.float64 if dtype == DataType.float64 else numpy.int32
             )
         if result.dtype != expected:
             raise TypeError(f"{prefix} must have dtype {expected}, not {result.dtype}")
diff --git a/madspace/src/compgraphs/function.cpp b/madspace/src/compgraphs/function.cpp
index 325f5fdf9c..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
@@ -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 543e4d0277..ee951db7fa 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(
@@ -941,7 +943,7 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const {
         case UMAMI_IN_RANDOM_DIAGRAM:
         case UMAMI_IN_MOMENTUM_FRACTION:
         case UMAMI_IN_SCALE:
-            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(
@@ -957,7 +959,8 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const {
         case UMAMI_IN_PARTICLE_COUNT:
         case UMAMI_IN_PDG_ID:
         case UMAMI_IN_BEAM_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
@@ -966,7 +969,8 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const {
             }
             break;
         case UMAMI_IN_PDG_IDS:
-            if (input_type.dtype != DataType::dt_int || input_type.shape.size() != 1) {
+            if (input_type.dtype != DataType::dt_int32 ||
+                input_type.shape.size() != 1) {
                 throw std::invalid_argument(
                     std::format(
                         "matrix_element, argument {}: expected array of integers", i + 1
@@ -999,17 +1003,17 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const {
         case UMAMI_OUT_FACT_SCALE_2:
         case UMAMI_OUT_PDF:
         case UMAMI_OUT_BIAS:
-            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/optimizer.cpp b/madspace/src/compgraphs/optimizer.cpp
index 3435dd3fb7..c43347665e 100644
--- a/madspace/src/compgraphs/optimizer.cpp
+++ b/madspace/src/compgraphs/optimizer.cpp
@@ -1,6 +1,7 @@
 #include "madspace/compgraphs/optimizer.hpp"
 
 #include <algorithm>
+#include <format>
 #include <numeric>
 #include <ranges>
 #include <unordered_map>
@@ -80,3 +81,109 @@ Function madspace::sort_breadth_first(const Function& function) {
     }
     return func_out;
 }
+
+Function madspace::fuse_instructions(const Function& function) {
+    Function func_out = function;
+    auto& instructions = func_out._instructions;
+    std::vector<int> producer(function.locals().size(), -1);
+    std::vector<int> use_count(function.locals().size(), 0);
+    for (std::size_t i = 0; i < instructions.size(); ++i) {
+        for (auto& output : instructions[i].outputs) {
+            producer.at(output.local_index) = i;
+        }
+        for (auto& input : instructions[i].inputs) {
+            if (input.local_index != -1) {
+                ++use_count.at(input.local_index);
+            }
+        }
+    }
+    for (auto& output : function.outputs()) {
+        ++use_count.at(output.local_index);
+    }
+
+    auto find_instruction = [](const std::string& name) {
+        auto find = instruction_set.find(name);
+        if (find == instruction_set.end()) {
+            throw std::runtime_error(std::format("Unknown instruction '{}'", name));
+        }
+        return find->second.get();
+    };
+
+    // the precision of each activation instruction that fuses into a matmul
+    static const auto activation_precisions = [] {
+        std::unordered_map<std::string, std::string> precisions;
+        for (std::string activation :
+             {"relu", "leaky_relu", "elu", "gelu", "sigmoid", "softplus"}) {
+            for (std::string precision : {"float32", "float64"}) {
+                precisions.emplace(activation + "_" + precision, precision);
+            }
+        }
+        return precisions;
+    }();
+
+    std::vector<bool> removed(instructions.size());
+    bool any_fused = false;
+    for (std::size_t i = 0; i < instructions.size(); ++i) {
+        auto& call = instructions[i];
+        const std::string& name = call.instruction->name();
+        auto find_precision = activation_precisions.find(name);
+        bool is_activation = find_precision != activation_precisions.end();
+        if (name != "convert_float32" && name != "convert_float64" && !is_activation) {
+            continue;
+        }
+        auto& input = call.inputs.at(0);
+        if (input.local_index == -1) {
+            continue;
+        }
+        int source_index = producer.at(input.local_index);
+        if (source_index == -1 || removed.at(source_index)) {
+            continue;
+        }
+        auto& source = instructions.at(source_index);
+        // the fused instruction takes the place of the producer
+        if (source.stream_index != call.stream_index) {
+            continue;
+        }
+        const std::string& source_name = source.instruction->name();
+        if (is_activation) {
+            if (source_name != "matmul_" + find_precision->second) {
+                continue;
+            }
+            // the matmul output stays an output, the activation needs it backward
+            source.instruction = find_instruction("matmul_" + name);
+            source.outputs.push_back(call.outputs.at(0));
+        } else if (name == "convert_float32") {
+            bool fusable = source_name == "select" || source_name == "rqs_forward" ||
+                source_name == "rqs_inverse";
+            if (!fusable || source.outputs.at(0).local_index != input.local_index ||
+                input.type.dtype != DataType::dt_float64) {
+                continue;
+            }
+            source.instruction = find_instruction(source_name + "_with_float32");
+            source.outputs.push_back(call.outputs.at(0));
+        } else {
+            if (source_name != "matmul_float32" ||
+                use_count.at(input.local_index) != 1) {
+                continue;
+            }
+            source.instruction = find_instruction("matmul_to_float64");
+            source.outputs = {call.outputs.at(0)};
+        }
+        producer.at(call.outputs.at(0).local_index) = source_index;
+        removed.at(i) = true;
+        any_fused = true;
+    }
+    if (!any_fused) {
+        return func_out;
+    }
+
+    std::vector<InstructionCall> kept;
+    kept.reserve(instructions.size());
+    for (std::size_t i = 0; i < instructions.size(); ++i) {
+        if (!removed.at(i)) {
+            kept.push_back(std::move(instructions[i]));
+        }
+    }
+    instructions = std::move(kept);
+    return func_out;
+}
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..8cd1f43da7 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,38 +248,60 @@ void op_matrix_element(
     );
 }
 
-template <typename D>
+// out[j * batch_size + i] = in[...] + bias[j] after a gemm with beta = 0, in the
+// precision F of the matmul. matmul and matmul_to_float64 share it to round alike.
+template <typename F, typename O>
+void add_bias(
+    const F* in, const F* bias, O* out, std::size_t batch_size, std::size_t dims_out
+) {
+    for (std::size_t j = 0; j < dims_out; ++j) {
+        for (std::size_t i = 0; i < batch_size; ++i) {
+            std::size_t k = j * batch_size + i;
+            out[k] = static_cast<O>(in[k] + bias[j]);
+        }
+    }
+}
+
+template <typename F, typename D>
 void op_matmul(
     const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device
 ) {
     auto input = locals[instruction.input_indices[0]].contiguous(device);
     auto weight = locals[instruction.input_indices[1]].contiguous(device);
-    auto bias_orig = locals[instruction.input_indices[2]];
-    auto bias = bias_orig.contiguous(device);
+    auto bias = locals[instruction.input_indices[2]].contiguous(device);
     auto& output = locals[instruction.output_indices[0]];
     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.copy_from(bias_orig, device);
+    output = Tensor(dtype_of<F>(), {batch_size, dims_out}, device);
     if (batch_size == 0) {
         return;
     }
 
     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);
+    device.submit([=, out = output]() mutable {
+        auto* out_ptr = static_cast<F*>(out.data());
+        auto* bias_ptr = static_cast<F*>(bias.data());
+        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,
+            0.,
+            out_ptr,
+            batch_size
+        );
+        add_bias(out_ptr, bias_ptr, out_ptr, batch_size, dims_out);
     });
 }
 
-template <typename D>
+template <typename F, typename D>
 void backward_op_matmul(
     const CpuRuntime::Instruction& instruction,
     TensorVec& locals,
@@ -220,18 +320,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 +336,332 @@ 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
+        );
+    });
+}
+
+// input, weight and bias gradients of a matmul from the gradient of its output, for
+// the backward passes that run in a single job
+template <typename F>
+void matmul_grads(
+    const F* output_grad,
+    const F* input,
+    const F* weight,
+    F* input_grad,
+    F* weight_grad,
+    F* bias_grad,
+    std::size_t batch_size,
+    std::size_t dims_in,
+    std::size_t dims_out
+) {
+    auto* grad = const_cast<F*>(output_grad);
+    // input_grad += output_grad * weight
+    gemm<F>(
+        'N',
+        'N',
+        batch_size,
+        dims_in,
+        dims_out,
+        1.,
+        grad,
+        batch_size,
+        const_cast<F*>(weight),
+        dims_out,
+        1.,
+        input_grad,
+        batch_size
+    );
+    // weight_grad += output_grad.T * input
+    gemm<F>(
+        'T',
+        'N',
+        dims_out,
+        dims_in,
+        batch_size,
+        1.,
+        grad,
+        batch_size,
+        const_cast<F*>(input),
+        batch_size,
+        1.,
+        weight_grad,
+        dims_out
+    );
+    // bias_grad += sum_i output_grad_ij
+    std::vector<F> ones(batch_size, 1.);
+    gemv<F>(
+        'T',
+        batch_size,
+        dims_out,
+        1.,
+        grad,
+        batch_size,
+        ones.data(),
+        1,
+        1.,
+        bias_grad,
+        1
+    );
+}
+
+// matmul with a float32 input and a float64 output. Everything runs in a single job,
+// because the CPU runtime offers only one barrier per instruction.
+template <typename D>
+void op_matmul_to_float64(
+    const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device
+) {
+    auto input = locals[instruction.input_indices[0]].contiguous(device);
+    auto weight = locals[instruction.input_indices[1]].contiguous(device);
+    auto bias = locals[instruction.input_indices[2]].contiguous(device);
+    auto& output = locals[instruction.output_indices[0]];
+    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_float64, {batch_size, dims_out}, device);
+    if (batch_size == 0) {
+        return;
+    }
+    Tensor tmp(DataType::dt_float32, {batch_size, dims_out}, device);
+
+    device.sync_barrier();
+    device.submit([=, out = output]() mutable {
+        auto* tmp_ptr = static_cast<float*>(tmp.data());
+        auto* bias_ptr = static_cast<float*>(bias.data());
+        gemm<float>(
+            'N',
+            'T',
+            batch_size,
+            dims_out,
+            dims_in,
+            1.,
+            static_cast<float*>(input.data()),
+            batch_size,
+            static_cast<float*>(weight.data()),
+            dims_out,
+            0.,
+            tmp_ptr,
+            batch_size
+        );
+        add_bias(
+            tmp_ptr, bias_ptr, static_cast<double*>(out.data()), batch_size, dims_out
+        );
+    });
+}
+
+template <typename D>
+void backward_op_matmul_to_float64(
+    const CpuRuntime::Instruction& instruction,
+    TensorVec& locals,
+    TensorVec& local_grads,
+    const D& device
+) {
+    using F = float;
+    auto input = locals[instruction.input_indices[0]].contiguous(device);
+    auto weight = locals[instruction.input_indices[1]].contiguous(device);
+    auto output_grad = local_grads[instruction.output_indices[0]].contiguous(device);
+    auto& input_grad = local_grads[instruction.input_indices[0]];
+    auto& weight_grad = local_grads[instruction.input_indices[1]];
+    auto& bias_grad = local_grads[instruction.input_indices[2]];
+    std::size_t batch_size = input.size(0);
+    std::size_t dims_in = input.size(1);
+    std::size_t dims_out = weight.size(1);
+
+    if (!input_grad) {
+        input_grad =
+            Tensor(dtype_of<F>(), input.shape(), device, AllocHint::local_grad);
+    }
+    if (!weight_grad) {
+        weight_grad =
+            Tensor(dtype_of<F>(), weight.shape(), device, AllocHint::local_grad);
+    }
+    if (!bias_grad) {
+        bias_grad = Tensor(dtype_of<F>(), {1, dims_out}, device, AllocHint::local_grad);
+    }
+    if (batch_size == 0) {
+        return;
+    }
+    Tensor grad_f(dtype_of<F>(), output_grad.shape(), device);
+
+    device.sync_barrier();
+    device.submit([=]() mutable {
+        auto* grad_f_ptr = static_cast<F*>(grad_f.data());
+        auto* grad_ptr = static_cast<double*>(output_grad.data());
+        for (std::size_t i = 0; i < batch_size * dims_out; ++i) {
+            grad_f_ptr[i] = static_cast<F>(grad_ptr[i]);
+        }
+        matmul_grads(
+            grad_f_ptr,
+            static_cast<F*>(input.data()),
+            static_cast<F*>(weight.data()),
+            static_cast<F*>(input_grad.data()),
+            static_cast<F*>(weight_grad.data()),
+            static_cast<F*>(bias_grad.data()),
+            batch_size,
+            dims_in,
+            dims_out
+        );
+    });
+}
+
+// matmul followed by the activation Act, with the outputs linear (before the
+// activation) and y. Everything runs in a single job, as in op_matmul_to_float64.
+template <typename Act, typename F, typename D>
+void op_matmul_activation(
+    const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device
+) {
+    using T = CpuTypesT<F>;
+    auto input = locals[instruction.input_indices[0]].contiguous(device);
+    auto weight = locals[instruction.input_indices[1]].contiguous(device);
+    auto bias = locals[instruction.input_indices[2]].contiguous(device);
+    auto& linear = locals[instruction.output_indices[0]];
+    auto& output = locals[instruction.output_indices[1]];
+    std::size_t batch_size = input.size(0);
+    std::size_t dims_in = input.size(1);
+    std::size_t dims_out = weight.size(1);
+    linear = Tensor(dtype_of<F>(), {batch_size, dims_out}, device);
+    output = Tensor(dtype_of<F>(), {batch_size, dims_out}, device);
+    if (batch_size == 0) {
+        return;
+    }
+
+    device.sync_barrier();
+    device.submit([=, lin = linear, out = output]() mutable {
+        auto* lin_ptr = static_cast<F*>(lin.data());
+        auto* out_ptr = static_cast<F*>(out.data());
+        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,
+            0.,
+            lin_ptr,
+            batch_size
+        );
+        add_bias(lin_ptr, static_cast<F*>(bias.data()), lin_ptr, batch_size, dims_out);
+        for (std::size_t k = 0; k < batch_size * dims_out; ++k) {
+            out_ptr[k] = Act::template forward<T>(lin_ptr[k]);
+        }
+    });
+}
+
+template <typename Act, typename F, typename D>
+void backward_op_matmul_activation(
+    const CpuRuntime::Instruction& instruction,
+    TensorVec& locals,
+    TensorVec& local_grads,
+    const D& device
+) {
+    using T = CpuTypesT<F>;
+    auto input = locals[instruction.input_indices[0]].contiguous(device);
+    auto weight = locals[instruction.input_indices[1]].contiguous(device);
+    auto linear = locals[instruction.output_indices[0]].contiguous(device);
+    auto output = locals[instruction.output_indices[1]].contiguous(device);
+    auto linear_grad = local_grads[instruction.output_indices[0]].contiguous(device);
+    auto output_grad = local_grads[instruction.output_indices[1]].contiguous(device);
+    auto& input_grad = local_grads[instruction.input_indices[0]];
+    auto& weight_grad = local_grads[instruction.input_indices[1]];
+    auto& bias_grad = local_grads[instruction.input_indices[2]];
+    std::size_t batch_size = input.size(0);
+    std::size_t dims_in = input.size(1);
+    std::size_t dims_out = weight.size(1);
+
+    if (!input_grad) {
+        input_grad =
+            Tensor(dtype_of<F>(), input.shape(), device, AllocHint::local_grad);
+    }
+    if (!weight_grad) {
+        weight_grad =
+            Tensor(dtype_of<F>(), weight.shape(), device, AllocHint::local_grad);
+    }
+    if (!bias_grad) {
+        bias_grad = Tensor(dtype_of<F>(), {1, dims_out}, device, AllocHint::local_grad);
+    }
+    if (batch_size == 0) {
+        return;
+    }
+    Tensor grad(dtype_of<F>(), {batch_size, dims_out}, device);
+
+    device.sync_barrier();
+    device.submit([=]() mutable {
+        auto* grad_ptr = static_cast<F*>(grad.data());
+        auto* lin_ptr = static_cast<F*>(linear.data());
+        auto* out_ptr = static_cast<F*>(output.data());
+        auto* lin_grad_ptr = static_cast<F*>(linear_grad.data());
+        auto* out_grad_ptr = static_cast<F*>(output_grad.data());
+        // the gradient of linear also gets the contributions of its other uses
+        for (std::size_t k = 0; k < batch_size * dims_out; ++k) {
+            grad_ptr[k] = lin_grad_ptr[k] +
+                Act::template backward<T>(lin_ptr[k], out_ptr[k], out_grad_ptr[k]);
+        }
+        matmul_grads(
+            static_cast<const F*>(grad_ptr),
+            static_cast<F*>(input.data()),
+            static_cast<F*>(weight.data()),
+            static_cast<F*>(input_grad.data()),
+            static_cast<F*>(weight_grad.data()),
+            static_cast<F*>(bias_grad.data()),
+            batch_size,
+            dims_in,
+            dims_out
+        );
     });
 }
 
@@ -287,7 +673,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 +718,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 +782,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 +848,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 +895,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 +939,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 +1060,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 +1095,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 +1124,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 +1146,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_int32, {batch_size}, device);
     auto flat_view = output.flat_view<me_int_t, 1>(0);
     auto& runtime = instruction.runtime;
     device.foreach (
@@ -788,8 +1174,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 +1232,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 +1301,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 +1360,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 +1520,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 +1583,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 +1602,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 +1776,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 +1795,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 +1967,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 +2006,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..30ef451224 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)) {};
@@ -47,45 +52,64 @@ struct BVec {
     __m256d v;
 };
 
-inline __m128i stride_seq(std::size_t stride) {
-    return _mm_mullo_epi32(_mm_set1_epi32(stride), _mm_set_epi32(3, 2, 1, 0));
-}
-
-inline __m128i
-mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) {
-    return _mm_add_epi32(
-        _mm_mullo_epi32(truncate_i64_to_i32(indices), _mm_set1_epi32(index_stride)),
-        stride_seq(batch_stride)
+// 64-bit element offsets index * index_stride + lane * batch_stride of four lanes,
+// from the low 32 bits of the 64-bit indices (indices >= 0, strides < 2^32)
+inline __m256i mem_indices(
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    __m256i indices,
+    long long first_lane = 0
+) {
+    return _mm256_add_epi64(
+        _mm256_mul_epu32(indices, _mm256_set1_epi64x(index_stride)),
+        _mm256_mul_epu32(
+            _mm256_set1_epi64x(batch_stride),
+            _mm256_add_epi64(
+                _mm256_set1_epi64x(first_lane), _mm256_set_epi64x(3, 2, 1, 0)
+            )
+        )
     );
 }
 
 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(
+    return _mm256_i64gather_pd(
         base_ptr, mem_indices(batch_stride, index_stride, indices), 8
     );
 }
 
 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(
+    return _mm256_cvtepi32_epi64(_mm256_i64gather_epi32(
         base_ptr, mem_indices(batch_stride, index_stride, indices), 4
     ));
 }
 
 // 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 +154,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 +239,244 @@ 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;
+};
+
+// 64-bit offsets of the lower (half = 0) or upper four of the eight lanes
+template <int half>
+inline __m256i
+mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) {
+    return mem_indices(
+        batch_stride,
+        index_stride,
+        _mm256_cvtepi32_epi64(_mm256_extracti128_si256(indices, half)),
+        4 * half
+    );
+}
+
+inline FVec32 vgather(
+    VecTag<FVec32>,
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    return _mm256_set_m128(
+        _mm256_i64gather_ps(
+            base_ptr, mem_indices32<1>(batch_stride, index_stride, indices), 4
+        ),
+        _mm256_i64gather_ps(
+            base_ptr, mem_indices32<0>(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_set_m128i(
+        _mm256_i64gather_epi32(
+            base_ptr, mem_indices32<1>(batch_stride, index_stride, indices), 4
+        ),
+        _mm256_i64gather_epi32(
+            base_ptr, mem_indices32<0>(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), indices);
+    for (int i = 0; i < 8; ++i) {
+        base_ptr[index_stride * indices_buf[i] + batch_stride * 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), indices);
+    for (int i = 0; i < 8; ++i) {
+        base_ptr[index_stride * indices_buf[i] + batch_stride * 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..3789ee7fd9 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,258 @@ 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;
+};
+
+// 64-bit offsets of the lower (half = 0) or upper four of the eight lanes
+template <int half>
+inline __m256i
+mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) {
+    return _mm256_add_epi64(
+        mem_indices(
+            batch_stride,
+            index_stride,
+            _mm256_cvtepi32_epi64(_mm256_extracti128_si256(indices, half))
+        ),
+        _mm256_set1_epi64x(4 * half * batch_stride)
+    );
+}
+
+inline FVec32 vgather(
+    VecTag<FVec32>,
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    return _mm256_set_m128(
+        _mm256_i64gather_ps(
+            base_ptr, mem_indices32<1>(batch_stride, index_stride, indices), 4
+        ),
+        _mm256_i64gather_ps(
+            base_ptr, mem_indices32<0>(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_set_m128i(
+        _mm256_i64gather_epi32(
+            base_ptr, mem_indices32<1>(batch_stride, index_stride, indices), 4
+        ),
+        _mm256_i64gather_epi32(
+            base_ptr, mem_indices32<0>(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_i64scatter_ps(
+        base_ptr,
+        mem_indices32<0>(batch_stride, index_stride, indices),
+        _mm256_castps256_ps128(values),
+        4
+    );
+    _mm256_i64scatter_ps(
+        base_ptr,
+        mem_indices32<1>(batch_stride, index_stride, indices),
+        _mm256_extractf128_ps(values, 1),
+        4
+    );
+}
+
+inline void vscatter(
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices,
+    IVec32 values
+) {
+    _mm256_i64scatter_epi32(
+        base_ptr,
+        mem_indices32<0>(batch_stride, index_stride, indices),
+        _mm256_castsi256_si128(values),
+        4
+    );
+    _mm256_i64scatter_epi32(
+        base_ptr,
+        mem_indices32<1>(batch_stride, index_stride, indices),
+        _mm256_extracti128_si256(values, 1),
+        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..1e7021f4c2 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) {};
@@ -38,51 +43,66 @@ struct BVec {
     __mmask8 v;
 };
 
-inline __m256i stride_seq(std::size_t stride) {
-    return _mm256_mullo_epi32(
-        _mm256_set1_epi32(stride), _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0)
-    );
-}
-
-inline __m256i
-mem_indices(std::size_t batch_stride, std::size_t index_stride, IVec indices) {
-    return _mm256_add_epi32(
-        _mm256_mullo_epi32(
-            _mm512_cvtepi64_epi32(indices), _mm256_set1_epi32(index_stride)
-        ),
-        stride_seq(batch_stride)
+// 64-bit element offsets index * index_stride + lane * batch_stride of eight lanes,
+// from the low 32 bits of the 64-bit indices (indices >= 0, strides < 2^32)
+inline __m512i mem_indices(
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    __m512i indices,
+    long long first_lane = 0
+) {
+    return _mm512_add_epi64(
+        _mm512_mul_epu32(indices, _mm512_set1_epi64(index_stride)),
+        _mm512_mul_epu32(
+            _mm512_set1_epi64(batch_stride),
+            _mm512_add_epi64(
+                _mm512_set1_epi64(first_lane), _mm512_set_epi64(7, 6, 5, 4, 3, 2, 1, 0)
+            )
+        )
     );
 }
 
 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(
+    return _mm512_i64gather_pd(
         mem_indices(batch_stride, index_stride, indices), base_ptr, 8
     );
 }
 
 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
+    return _mm512_cvtepi32_epi64(_mm512_i64gather_epi32(
+        mem_indices(batch_stride, index_stride, indices), base_ptr, 4
     ));
 }
 
 // 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,
@@ -91,7 +111,7 @@ inline void vscatter(
     IVec indices,
     FVec values
 ) {
-    _mm512_i32scatter_pd(
+    _mm512_i64scatter_pd(
         base_ptr, mem_indices(batch_stride, index_stride, indices), values, 8
     );
 }
@@ -119,7 +139,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 +242,271 @@ 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;
+};
+
+// 64-bit offsets of the lower (half = 0) or upper eight of the sixteen lanes
+template <int half>
+inline __m512i
+mem_indices32(std::size_t batch_stride, std::size_t index_stride, IVec32 indices) {
+    return mem_indices(
+        batch_stride,
+        index_stride,
+        _mm512_cvtepi32_epi64(_mm512_extracti64x4_epi64(indices, half)),
+        8 * half
+    );
+}
+
+// the halves of a float32 vector (no AVX512DQ for the 32-bit inserts and extracts)
+inline __m512 join_halves(__m256 low, __m256 high) {
+    return _mm512_castpd_ps(_mm512_insertf64x4(
+        _mm512_castps_pd(_mm512_castps256_ps512(low)), _mm256_castps_pd(high), 1
+    ));
+}
+
+template <int half>
+inline __m256 get_half(__m512 v) {
+    return _mm256_castpd_ps(_mm512_extractf64x4_pd(_mm512_castps_pd(v), half));
+}
+
+inline FVec32 vgather(
+    VecTag<FVec32>,
+    float* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices
+) {
+    return join_halves(
+        _mm512_i64gather_ps(
+            mem_indices32<0>(batch_stride, index_stride, indices), base_ptr, 4
+        ),
+        _mm512_i64gather_ps(
+            mem_indices32<1>(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_inserti64x4(
+        _mm512_castsi256_si512(_mm512_i64gather_epi32(
+            mem_indices32<0>(batch_stride, index_stride, indices), base_ptr, 4
+        )),
+        _mm512_i64gather_epi32(
+            mem_indices32<1>(batch_stride, index_stride, indices), base_ptr, 4
+        ),
+        1
+    );
+}
+
+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_i64scatter_ps(
+        base_ptr,
+        mem_indices32<0>(batch_stride, index_stride, indices),
+        get_half<0>(values),
+        4
+    );
+    _mm512_i64scatter_ps(
+        base_ptr,
+        mem_indices32<1>(batch_stride, index_stride, indices),
+        get_half<1>(values),
+        4
+    );
+}
+
+inline void vscatter(
+    int* base_ptr,
+    std::size_t batch_stride,
+    std::size_t index_stride,
+    IVec32 indices,
+    IVec32 values
+) {
+    _mm512_i64scatter_epi32(
+        base_ptr,
+        mem_indices32<0>(batch_stride, index_stride, indices),
+        _mm512_extracti64x4_epi64(values, 0),
+        4
+    );
+    _mm512_i64scatter_epi32(
+        base_ptr,
+        mem_indices32<1>(batch_stride, index_stride, indices),
+        _mm512_extracti64x4_epi64(values, 1),
+        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 398fe70cc5..d6c6147131 100644
--- a/madspace/src/driver/context.cpp
+++ b/madspace/src/driver/context.cpp
@@ -387,7 +387,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 8e3bf523dc..01d2f9631d 100644
--- a/madspace/src/driver/systematics.cpp
+++ b/madspace/src/driver/systematics.cpp
@@ -125,7 +125,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)
+        );
     }
 }
 
@@ -214,7 +216,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());
@@ -266,13 +270,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),
@@ -293,32 +299,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;
@@ -409,11 +422,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});
     }
@@ -550,9 +562,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>();
@@ -582,7 +594,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];
@@ -600,8 +612,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);
@@ -623,9 +636,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>();
@@ -658,8 +671,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");
@@ -812,7 +826,8 @@ void SystematicsCalculator::compute(
         }
         if (has_ladder && use_pdf) {
             std::array<double, 2> lead_q{
-                buffer.particle(i, 0).ladder_lead_q(), buffer.particle(i, 1).ladder_lead_q()
+                buffer.particle(i, 0).ladder_lead_q(),
+                buffer.particle(i, 1).ladder_lead_q()
             };
             // an event of a channel without the ladder carries zeros
             in.has_ladder = lead_q[0] > 0. || lead_q[1] > 0.;
@@ -849,35 +864,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) {
@@ -894,9 +911,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) {
@@ -913,46 +929,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++];
+                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;
                     }
-                    if (has_beam2) {
-                        product *= values[n++];
-                    }
-                    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) {
@@ -1048,13 +1065,14 @@ void SystematicsCalculator::compute(
         auto nominal = ladder_weights(_nominal_pdf.value(), central);
         for (std::size_t k = 0; k < var_count; ++k) {
             auto& var = _variations.at(k);
-            bool varies_pdf = var.pdf_index >= 0 || var.muf != 1. || var.dyn != -1 ||
-                var.alps != 1.;
+            bool varies_pdf =
+                var.pdf_index >= 0 || var.muf != 1. || var.dyn != -1 || var.alps != 1.;
             if (!varies_pdf) {
                 continue;
             }
-            auto& evaluator = var.pdf_index >= 0 ? _member_pdfs.at(var.pdf_index)
-                                                 : _nominal_pdf.value();
+            auto& evaluator = var.pdf_index >= 0
+                ? _member_pdfs.at(var.pdf_index)
+                : _nominal_pdf.value();
             auto varied = ladder_weights(evaluator, var);
             for (std::size_t n = 0; n < ladder_events.size(); ++n) {
                 r_pdf[ladder_events[n]][k] =
@@ -1126,8 +1144,9 @@ void SystematicsCalculator::compute(
         }
         // an emission scale variation alone leaves the mu_R powers as they are
         bool varies_mur = key.mur != 1. || key.dyn != -1 || key.alpha_s_index != 0;
-        auto varied = varies_mur ? evaluate_alpha_s(key.alpha_s_index, q)
-                                 : std::vector<double>(count, 0.);
+        auto varied = varies_mur
+            ? evaluate_alpha_s(key.alpha_s_index, q)
+            : std::vector<double>(count, 0.);
         std::vector<double> ratio(count, 1.);
         if (varies_mur) {
             for (std::size_t i = 0; i < count; ++i) {
@@ -1348,8 +1367,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..402fc03d73 100644
--- a/madspace/src/gpu/runtime.cu
+++ b/madspace/src/gpu/runtime.cu
@@ -2,9 +2,11 @@
 
 #include <algorithm>
 #include <array>
+#include <cstdlib>
 #include <format>
 #include <functional>
 #include <tuple>
+#include <type_traits>
 
 #include <thrust/copy.h>
 #include <thrust/device_ptr.h>
@@ -125,6 +127,44 @@ 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...);
+    }
+}
+
+// out[j * batch_size + i] += bias[j], over the output of a gemm
+template <typename F>
+__global__ void
+kernel_add_bias(std::size_t total, std::size_t batch_size, const F* bias, F* out) {
+    std::size_t i = blockDim.x * static_cast<std::size_t>(blockIdx.x) + threadIdx.x;
+    if (i < total) {
+        out[i] += bias[i / batch_size];
+    }
+}
+
+// out[j * batch_size + i] = double(tmp[...] + bias[j]): bias, convert to float64 in
+// one pass, rounding exactly like kernel_add_bias followed by a conversion
+template <typename F>
+__global__ void kernel_bias_to_double(
+    std::size_t total,
+    std::size_t batch_size,
+    const F* __restrict__ tmp,
+    const F* __restrict__ bias,
+    double* __restrict__ out
+) {
+    std::size_t i = blockDim.x * static_cast<std::size_t>(blockIdx.x) + threadIdx.x;
+    if (i < total) {
+        out[i] = static_cast<double>(tmp[i] + bias[i / batch_size]);
+    }
+}
+
+// the bias is added after the gemm (beta = 0), so that matmul_to_float64 rounds
+// exactly like matmul followed by convert_float64
+template <typename F>
 void op_matmul(
     const GpuRuntime::Instruction& instruction,
     TensorVec& locals,
@@ -141,49 +181,284 @@ void op_matmul(
     std::size_t dims_in = input.size(1);
     std::size_t dims_out = weight.size(1);
     output = Tensor(
-        DataType::dt_float,
+        dtype_of<F>(), {batch_size, dims_out}, device, instruction.output_alloc_hints[0]
+    );
+    if (batch_size == 0) {
+        return;
+    }
+
+    gpublasHandle_t handle = instruction.runtime.gpublas_handle();
+    check_error(gpublasSetStream(handle, device.stream()));
+    F alpha = 1., beta = 0.;
+    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
+        )
+    );
+    std::size_t total = batch_size * dims_out;
+    launch_kernel(
+        kernel_add_bias<F>,
+        total,
+        device.stream(),
+        total,
+        batch_size,
+        static_cast<const F*>(bias.data()),
+        static_cast<F*>(output.data())
+    );
+    input.reset(device);
+    weight.reset(device);
+    bias.reset(device);
+}
+
+// matmul of float32 inputs with the float32 -> float64 conversion of the result
+// fused into the bias epilogue (three launches become two)
+void op_matmul_to_float64(
+    const GpuRuntime::Instruction& instruction,
+    TensorVec& locals,
+    const AsyncGpuDevice& device
+) {
+    using F = float;
+    auto input =
+        locals[instruction.input_indices[0]].contiguous(device, AllocHint::temporary);
+    auto weight =
+        locals[instruction.input_indices[1]].contiguous(device, AllocHint::temporary);
+    auto bias =
+        locals[instruction.input_indices[2]].contiguous(device, AllocHint::temporary);
+    auto& output = locals[instruction.output_indices[0]];
+    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_float64,
         {batch_size, dims_out},
         device,
         instruction.output_alloc_hints[0]
     );
-    output.copy_from(bias, device);
     if (batch_size == 0) {
         return;
     }
 
+    Tensor tmp(dtype_of<F>(), {batch_size, dims_out}, device, AllocHint::temporary);
     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,
+    F alpha = 1., beta = 0.;
+    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*>(tmp.data()),
+            batch_size
+        )
+    );
+    std::size_t total = batch_size * dims_out;
+    launch_kernel(
+        kernel_bias_to_double<F>,
+        total,
+        device.stream(),
+        total,
         batch_size,
-        dims_out,
-        dims_in,
-        &alpha,
-        static_cast<double*>(input.data()),
+        static_cast<const F*>(tmp.data()),
+        static_cast<const F*>(bias.data()),
+        static_cast<double*>(output.data())
+    );
+    tmp.reset(device);
+    input.reset(device);
+    weight.reset(device);
+    bias.reset(device);
+}
+
+// linear[j * batch_size + i] += bias[j] and y = Act(linear) in one pass over the
+// output of a gemm, rounding exactly like kernel_add_bias and the activation kernel
+template <typename Act, typename F>
+__global__ void kernel_bias_activation(
+    std::size_t total,
+    std::size_t batch_size,
+    const F* __restrict__ bias,
+    F* __restrict__ linear,
+    F* __restrict__ y
+) {
+    std::size_t i = blockDim.x * static_cast<std::size_t>(blockIdx.x) + threadIdx.x;
+    if (i < total) {
+        F value = linear[i] + bias[i / batch_size];
+        linear[i] = value;
+        y[i] = Act::template forward<GpuTypesT<F>>(value);
+    }
+}
+
+// matmul followed by the activation Act, with the outputs linear (before the
+// activation, for the backward pass) and y: two launches instead of three
+template <typename Act, typename F>
+void op_matmul_activation(
+    const GpuRuntime::Instruction& instruction,
+    TensorVec& locals,
+    const AsyncGpuDevice& device
+) {
+    auto input =
+        locals[instruction.input_indices[0]].contiguous(device, AllocHint::temporary);
+    auto weight =
+        locals[instruction.input_indices[1]].contiguous(device, AllocHint::temporary);
+    auto bias =
+        locals[instruction.input_indices[2]].contiguous(device, AllocHint::temporary);
+    auto& linear = locals[instruction.output_indices[0]];
+    auto& output = locals[instruction.output_indices[1]];
+    std::size_t batch_size = input.size(0);
+    std::size_t dims_in = input.size(1);
+    std::size_t dims_out = weight.size(1);
+    linear = Tensor(
+        dtype_of<F>(), {batch_size, dims_out}, device, instruction.output_alloc_hints[0]
+    );
+    output = Tensor(
+        dtype_of<F>(), {batch_size, dims_out}, device, instruction.output_alloc_hints[1]
+    );
+    if (batch_size == 0) {
+        return;
+    }
+
+    gpublasHandle_t handle = instruction.runtime.gpublas_handle();
+    check_error(gpublasSetStream(handle, device.stream()));
+    F alpha = 1., beta = 0.;
+    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*>(linear.data()),
+            batch_size
+        )
+    );
+    std::size_t total = batch_size * dims_out;
+    launch_kernel(
+        kernel_bias_activation<Act, F>,
+        total,
+        device.stream(),
+        total,
         batch_size,
-        static_cast<double*>(weight.data()),
-        dims_out,
-        &beta,
-        static_cast<double*>(output.data()),
-        batch_size
-    ));
+        static_cast<const F*>(bias.data()),
+        static_cast<F*>(linear.data()),
+        static_cast<F*>(output.data())
+    );
     input.reset(device);
     weight.reset(device);
     bias.reset(device);
 }
 
-__global__ void
-kernel_one(std::size_t batch_size, GpuTensorView<double, 1, true> output) {
-    me_int_t i = blockDim.x * blockIdx.x + threadIdx.x;
-    if (i < batch_size) {
-        output[i] = 1.;
+// bias_grad[j] += sum_i output_grad[i, j]; the matrix is column-major with
+// leading dimension batch_size, so each column is a contiguous segment. One block
+// per column: single launch, no atomics (deterministic) and no temporary storage.
+// Accumulates in double, as the float32 cuBLAS paths would otherwise split the
+// reduction over several kernels.
+constexpr int BIAS_GRAD_THREADS = 512;
+
+template <typename F>
+__global__ void kernel_bias_grad(
+    std::size_t batch_size, const F* __restrict__ output_grad, F* bias_grad
+) {
+    using BlockReduce = cub::BlockReduce<double, BIAS_GRAD_THREADS>;
+    __shared__ typename BlockReduce::TempStorage temp_storage;
+    const F* column = output_grad + blockIdx.x * batch_size;
+    double sum = 0.;
+    for (std::size_t i = threadIdx.x; i < batch_size; i += BIAS_GRAD_THREADS) {
+        sum += column[i];
+    }
+    sum = BlockReduce(temp_storage).Sum(sum);
+    if (threadIdx.x == 0) {
+        bias_grad[blockIdx.x] += static_cast<F>(sum);
     }
 }
 
-void backward_op_matmul(
+// the same reduction for a float64 output gradient, which is also converted to F
+// for the gemms: one read of the gradient instead of a separate convert kernel
+template <typename F>
+__global__ void kernel_bias_grad_convert(
+    std::size_t batch_size,
+    const double* __restrict__ output_grad,
+    F* __restrict__ output_grad_f,
+    F* bias_grad
+) {
+    using BlockReduce = cub::BlockReduce<double, BIAS_GRAD_THREADS>;
+    __shared__ typename BlockReduce::TempStorage temp_storage;
+    std::size_t offset = blockIdx.x * batch_size;
+    double sum = 0.;
+    for (std::size_t i = threadIdx.x; i < batch_size; i += BIAS_GRAD_THREADS) {
+        F value = static_cast<F>(output_grad[offset + i]);
+        output_grad_f[offset + i] = value;
+        sum += value;
+    }
+    sum = BlockReduce(temp_storage).Sum(sum);
+    if (threadIdx.x == 0) {
+        bias_grad[blockIdx.x] += static_cast<F>(sum);
+    }
+}
+
+// the same reduction over the gradient of the input of the activation Act of
+// matmul_<activation>, linear_grad += Act'(linear, y) * y_grad, which is written
+// back for the gemms. It rounds like the activation backward and kernel_bias_grad.
+template <typename Act, typename F>
+__global__ void kernel_bias_grad_activation(
+    std::size_t batch_size,
+    const F* __restrict__ linear,
+    const F* __restrict__ y,
+    const F* __restrict__ y_grad,
+    F* __restrict__ linear_grad,
+    F* bias_grad
+) {
+    using BlockReduce = cub::BlockReduce<double, BIAS_GRAD_THREADS>;
+    __shared__ typename BlockReduce::TempStorage temp_storage;
+    std::size_t offset = blockIdx.x * batch_size;
+    double sum = 0.;
+    for (std::size_t i = threadIdx.x; i < batch_size; i += BIAS_GRAD_THREADS) {
+        std::size_t k = offset + i;
+        F value = linear_grad[k] +
+            Act::template backward<GpuTypesT<F>>(linear[k], y[k], y_grad[k]);
+        linear_grad[k] = value;
+        sum += value;
+    }
+    sum = BlockReduce(temp_storage).Sum(sum);
+    if (threadIdx.x == 0) {
+        bias_grad[blockIdx.x] += static_cast<F>(sum);
+    }
+}
+
+// How the gradient of the matmul output is obtained: as is (matmul), converted
+// from float64 (matmul_to_float64) or through an activation struct (matmul_<act>),
+// in the last two cases together with the bias gradient reduction
+struct OutputGrad {};
+struct OutputGradFromFloat64 {};
+
+template <typename F, typename Grad>
+void backward_matmul_impl(
     const GpuRuntime::Instruction& instruction,
     TensorVec& locals,
     TensorVec& local_grads,
@@ -205,30 +480,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,72 +500,136 @@ 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));
 
+    if constexpr (std::is_same_v<Grad, OutputGradFromFloat64>) {
+        // compute bias_grad += sum_i output_grad_ij and convert the gradient
+        Tensor output_grad_f(
+            dtype_of<F>(), output_grad.shape(), device, AllocHint::temporary
+        );
+        kernel_bias_grad_convert<F><<<dims_out, BIAS_GRAD_THREADS, 0, stream>>>(
+            batch_size,
+            static_cast<const double*>(output_grad.data()),
+            static_cast<F*>(output_grad_f.data()),
+            static_cast<F*>(bias_grad.data())
+        );
+        check_error();
+        output_grad.reset(device);
+        output_grad = output_grad_f;
+    } else if constexpr (!std::is_same_v<Grad, OutputGrad>) {
+        // the gradient of linear (output_grad) and bias_grad through the activation
+        auto linear = locals[instruction.output_indices[0]].contiguous(
+            device, AllocHint::temporary
+        );
+        auto y = locals[instruction.output_indices[1]].contiguous(
+            device, AllocHint::temporary
+        );
+        auto y_grad = local_grads[instruction.output_indices[1]].contiguous(
+            device, AllocHint::temporary
+        );
+        kernel_bias_grad_activation<Grad, F>
+            <<<dims_out, BIAS_GRAD_THREADS, 0, stream>>>(
+                batch_size,
+                static_cast<const F*>(linear.data()),
+                static_cast<const F*>(y.data()),
+                static_cast<const F*>(y_grad.data()),
+                static_cast<F*>(output_grad.data()),
+                static_cast<F*>(bias_grad.data())
+            );
+        check_error();
+        linear.reset(device);
+        y.reset(device);
+        y_grad.reset(device);
+    }
+
     // 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
-    ));
-
-    // compute bias_grad += sum_i output_grad_ij
-    Tensor ones(DataType::dt_float, {batch_size}, device, AllocHint::temporary);
-    launch_kernel(
-        kernel_one, batch_size, device.stream(), batch_size, ones.view<double, 1>()
+    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
+        )
     );
-    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);
+
+    if constexpr (std::is_same_v<Grad, OutputGrad>) {
+        // compute bias_grad += sum_i output_grad_ij
+        kernel_bias_grad<F><<<dims_out, BIAS_GRAD_THREADS, 0, stream>>>(
+            batch_size,
+            static_cast<const F*>(output_grad.data()),
+            static_cast<F*>(bias_grad.data())
+        );
+        check_error();
+    }
     input.reset(device);
     weight.reset(device);
     output_grad.reset(device);
 }
 
+template <typename F>
+void backward_op_matmul(
+    const GpuRuntime::Instruction& instruction,
+    TensorVec& locals,
+    TensorVec& local_grads,
+    const AsyncGpuDevice& device
+) {
+    backward_matmul_impl<F, OutputGrad>(instruction, locals, local_grads, device);
+}
+
+void backward_op_matmul_to_float64(
+    const GpuRuntime::Instruction& instruction,
+    TensorVec& locals,
+    TensorVec& local_grads,
+    const AsyncGpuDevice& device
+) {
+    backward_matmul_impl<float, OutputGradFromFloat64>(
+        instruction, locals, local_grads, device
+    );
+}
+
+template <typename Act, typename F>
+void backward_op_matmul_activation(
+    const GpuRuntime::Instruction& instruction,
+    TensorVec& locals,
+    TensorVec& local_grads,
+    const AsyncGpuDevice& device
+) {
+    backward_matmul_impl<F, Act>(instruction, locals, local_grads, device);
+}
+
 struct NotMinusOne {
     __device__ bool operator()(me_int_t val) { return val != -1; }
 };
@@ -326,9 +653,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 +724,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 +735,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 +810,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 +821,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 +885,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 +909,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 +932,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 +959,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 +1045,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 +1058,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 +1104,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 +1112,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 +1129,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 +1242,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 +1291,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 +1305,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 +1322,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 +1356,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 +1395,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,
@@ -1107,15 +1444,16 @@ void op_unweight(
 
     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]);
+            Tensor(DataType::dt_int32, {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 +1467,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 +1483,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 +1518,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 +1591,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 +1601,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 +1655,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 +1739,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,
@@ -1577,7 +1922,13 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) :
     // check_error(cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold,
     // &thresh));
 
-    Function function = sort_breadth_first(function_arg);
+    // MADSPACE_NO_FUSION disables the lowering of precision conversions, to compare
+    // the fused and the unfused instructions
+    Function function = sort_breadth_first(
+        std::getenv("MADSPACE_NO_FUSION")
+            ? function_arg
+            : fuse_instructions(function_arg)
+    );
 
     _locals_init.resize(function.locals().size());
     _requires_grad_init.resize(function.locals().size());
@@ -1794,6 +2145,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 +2499,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 +2526,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/gpu/tensor.cuh b/madspace/src/gpu/tensor.cuh
index 2bbfd73d10..3b52e60948 100644
--- a/madspace/src/gpu/tensor.cuh
+++ b/madspace/src/gpu/tensor.cuh
@@ -5,6 +5,7 @@
 #include "madspace/driver/tensor.hpp"
 
 #include <algorithm>
+#include <type_traits>
 
 namespace madspace {
 namespace gpu {
@@ -19,9 +20,9 @@ public:
     using SizeType = std::conditional_t<packed, std::size_t[_dim], std::size_t*>;
     static const int dim = _dim;
 
-    GpuTensorView(const TensorView<T, _dim>& view)
+    GpuTensorView(const TensorView<T, _dim>& view, bool atomic = false)
         requires(packed)
-        : _data(view.data()) {
+        : _data(view.data()), _atomic(atomic) {
         if constexpr (_dim > 0) {
             std::copy(view.shape(), view.shape() + _dim, std::begin(this->_shape));
             std::copy(view.stride(), view.stride() + _dim, std::begin(this->_stride));
@@ -32,15 +33,20 @@ public:
 
     __host__ __device__ GpuTensorView(const GpuTensorView<T, _dim, true>& view)
         requires(!packed)
-        : _data(view._data), _stride(&view._stride), _shape(&view._shape) {}
+        :
+        _data(view._data),
+        _stride(&view._stride),
+        _shape(&view._shape),
+        _atomic(view._atomic) {}
 
     GpuTensorView<T, _dim, packed>& operator=(GpuTensorView<T, _dim, packed>& value)
         requires(packed)
     = default;
 
-    __device__ GpuTensorView(T* data, std::size_t* stride, std::size_t* shape)
+    __device__
+    GpuTensorView(T* data, std::size_t* stride, std::size_t* shape, bool atomic = false)
         requires(!packed)
-        : _data(data), _stride(stride), _shape(shape) {}
+        : _data(data), _stride(stride), _shape(shape), _atomic(atomic) {}
 
     __device__ GpuTensorView(T& value)
         requires(!packed)
@@ -52,13 +58,13 @@ public:
     __device__ const GpuTensorView<T, _dim - 1> operator[](std::size_t index) const
         requires(_dim != 0)
     {
-        return {&_data[index * _stride[0]], &_stride[1], &_shape[1]};
+        return {&_data[index * _stride[0]], &_stride[1], &_shape[1], _atomic};
     }
 
     __device__ GpuTensorView<T, _dim - 1> operator[](std::size_t index)
         requires(_dim != 0)
     {
-        return {&_data[index * _stride[0]], &_stride[1], &_shape[1]};
+        return {&_data[index * _stride[0]], &_stride[1], &_shape[1], _atomic};
     }
 
     template <typename... I>
@@ -68,7 +74,7 @@ public:
         T* ptr = _data;
         int i = 0;
         ((ptr = &ptr[index * _stride[i++]]), ...);
-        return {ptr, &_stride[sizeof...(I)], &_shape[sizeof...(I)]};
+        return {ptr, &_stride[sizeof...(I)], &_shape[sizeof...(I)], _atomic};
     }
 
     template <typename... I>
@@ -78,7 +84,7 @@ public:
         T* ptr = _data;
         int i = 0;
         ((ptr = &ptr[index * _stride[i++]]), ...);
-        return {ptr, &_stride[sizeof...(I)], &_shape[sizeof...(I)]};
+        return {ptr, &_stride[sizeof...(I)], &_shape[sizeof...(I)], _atomic};
     }
 
     __device__ operator T() const
@@ -94,9 +100,17 @@ public:
         return value;
     }
 
+    // atomic only if the launch found a zero stride (broadcast gradient), where many
+    // threads accumulate into the same element
     __device__ T operator+=(T value)
         requires(_dim == 0)
     {
+        if constexpr (std::is_floating_point_v<T>) {
+            if (_atomic) {
+                atomicAdd(_data, value);
+                return value;
+            }
+        }
         *_data += value;
         return value;
     }
@@ -123,6 +137,7 @@ private:
     T* _data;
     SizeType _stride;
     SizeType _shape;
+    bool _atomic = false;
 };
 
 template <typename T>
@@ -151,14 +166,23 @@ template <typename F, int dims>
 struct get_views;
 template <typename... TParam, int dims>
 struct get_views<void (*)(TParam...), dims> {
+    std::size_t batch_size;
+
     template <typename... TArg>
     auto operator()(TArg&&... args) {
         return std::make_tuple([&]() {
             if constexpr (ScalarType<TArg>) {
                 return GpuScalarView(args);
             } else {
+                auto view =
+                    args->template view<typename TParam::DType, TParam::dim + dims>();
+                // a zero stride in a dimension spread over threads aliases elements
+                bool atomic = false;
+                for (int i = 0; i < dims; ++i) {
+                    atomic |= view.stride()[i] == 0 && (i > 0 || batch_size > 1);
+                }
                 return GpuTensorView<typename TParam::DType, TParam::dim + dims, true>(
-                    args->template view<typename TParam::DType, TParam::dim + dims>()
+                    view, atomic
                 );
             }
         }()...);
@@ -258,7 +282,7 @@ void tensor_foreach(
     }
     // get views to the tensors with the correct types based on the signature of func
     auto views = std::apply(
-        get_views<decltype(func), dims>(),
+        get_views<decltype(func), dims>{batch_size},
         std::tuple_cat(inputs, outputs, std::make_tuple(scalar_args...))
     );
     auto& first_view = std::get<0>(views);
diff --git a/madspace/src/kernels/nn.hpp b/madspace/src/kernels/nn.hpp
index 176329a5e8..8dc1bb758d 100644
--- a/madspace/src/kernels/nn.hpp
+++ b/madspace/src/kernels/nn.hpp
@@ -8,102 +8,219 @@ 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.
+//
+// Each activation is a struct with the value functions forward(x) and
+// backward(x, y, g), the latter the input gradient for the input x, the output y
+// and the output gradient g. The element-wise kernels and the matmuls with a fused
+// activation share them, so that both round alike.
+
+struct ReluActivation {
+    template <typename T>
+    static KERNELSPEC FVal<T> forward(FVal<T> x) {
+        using F = FVal<T>;
+        return where(x < F(0.), F(0.), x);
+    }
+    template <typename T>
+    static KERNELSPEC FVal<T> backward(FVal<T> x, FVal<T> y, FVal<T> g) {
+        using F = FVal<T>;
+        return where(x < F(0.), F(0.), g);
+    }
+};
+
+struct LeakyReluActivation {
+    template <typename T>
+    static KERNELSPEC FVal<T> forward(FVal<T> x) {
+        using F = FVal<T>;
+        return where(x < F(0.), x * F(0.01), x);
+    }
+    template <typename T>
+    static KERNELSPEC FVal<T> backward(FVal<T> x, FVal<T> y, FVal<T> g) {
+        using F = FVal<T>;
+        return where(x < F(0.), g * F(0.01), g);
+    }
+};
+
+struct EluActivation {
+    template <typename T>
+    static KERNELSPEC FVal<T> forward(FVal<T> x) {
+        using F = FVal<T>;
+        return where(x < F(0.), expm1(x), x);
+    }
+    template <typename T>
+    static KERNELSPEC FVal<T> backward(FVal<T> x, FVal<T> y, FVal<T> g) {
+        using F = FVal<T>;
+        return where(x < F(0.), g * exp(x), g);
+    }
+};
+
+struct GeluActivation {
+    template <typename T>
+    static KERNELSPEC FVal<T> forward(FVal<T> x) {
+        using F = FVal<T>;
+        return F(0.5) * x * (F(1.) + erf(F(SQRT_HALF) * x));
+    }
+    template <typename T>
+    static KERNELSPEC FVal<T> backward(FVal<T> x, FVal<T> y, FVal<T> g) {
+        using F = FVal<T>;
+        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);
+        return g * (cdf + x * pdf);
+    }
+};
+
+// the backward pass uses the output
+struct SigmoidActivation {
+    template <typename T>
+    static KERNELSPEC FVal<T> forward(FVal<T> x) {
+        using F = FVal<T>;
+        return F(1.) / (F(1.) + exp(-x));
+    }
+    template <typename T>
+    static KERNELSPEC FVal<T> backward(FVal<T> x, FVal<T> y, FVal<T> g) {
+        using F = FVal<T>;
+        return g * y * (F(1.) - y);
+    }
+};
+
+struct SoftplusActivation {
+    template <typename T>
+    static KERNELSPEC FVal<T> forward(FVal<T> x) {
+        using F = FVal<T>;
+        return where(x < F(20.), log1p(exp(x)), x);
+    }
+    template <typename T>
+    static KERNELSPEC FVal<T> backward(FVal<T> x, FVal<T> y, FVal<T> g) {
+        using F = FVal<T>;
+        F z = exp(x);
+        return where(x < F(20.), g * z / (F(1.) + z), g);
+    }
+};
+
 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);
+    output = ReluActivation::forward<T>(FVal<T>(input));
 }
 
 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);
+    FVal<T> x(input);
+    input_grad += ReluActivation::backward<T>(x, x, FVal<T>(output_grad));
 }
 
 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);
+    output = LeakyReluActivation::forward<T>(FVal<T>(input));
 }
 
 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);
+    FVal<T> x(input);
+    input_grad += LeakyReluActivation::backward<T>(x, x, FVal<T>(output_grad));
 }
 
 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);
+    output = EluActivation::forward<T>(FVal<T>(input));
 }
 
 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);
+    FVal<T> x(input);
+    input_grad += EluActivation::backward<T>(x, x, FVal<T>(output_grad));
 }
 
 template <typename T>
 KERNELSPEC void kernel_gelu(FIn<T, 0> input, FOut<T, 0> output) {
-    output = 0.5 * input * (1. + erf(SQRT_HALF * input));
+    output = GeluActivation::forward<T>(FVal<T>(input));
 }
 
 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);
+    FVal<T> x(input);
+    input_grad += GeluActivation::backward<T>(x, x, FVal<T>(output_grad));
 }
 
 template <typename T>
 KERNELSPEC void kernel_sigmoid(FIn<T, 0> input, FOut<T, 0> output) {
-    FVal<T> x = input;
-    output = 1. / (1. + exp(-x));
+    output = SigmoidActivation::forward<T>(FVal<T>(input));
 }
 
 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);
+    FVal<T> y(output);
+    input_grad += SigmoidActivation::backward<T>(y, y, FVal<T>(output_grad));
 }
 
 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);
+    output = SoftplusActivation::forward<T>(FVal<T>(input));
 }
 
 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);
+    FVal<T> x(input);
+    input_grad += SoftplusActivation::backward<T>(x, x, FVal<T>(output_grad));
 }
 
 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>
@@ -203,9 +320,13 @@ KERNELSPEC void backward_kernel_rqs_find_bin(
 }
 
 template <typename T>
-KERNELSPEC void kernel_rqs_forward(
-    FIn<T, 0> input, FIn<T, 1> condition, FOut<T, 0> output, FOut<T, 0> det
-) {
+struct RqsValues {
+    FVal<T> output;
+    FVal<T> det;
+};
+
+template <typename T>
+KERNELSPEC RqsValues<T> rqs_forward_values(FIn<T, 0> input, FIn<T, 1> condition) {
     FVal<T> width(condition[0]), height(condition[1]);
     FVal<T> cumwidth(condition[2]), cumheight(condition[3]);
     FVal<T> derivative_unorm(condition[4]), derivative_plus_one_unorm(condition[5]);
@@ -240,20 +361,46 @@ KERNELSPEC void kernel_rqs_forward(
     auto denominator =
         delta + (derivative + derivative_plus_one - two_delta) * theta_one_minus_theta;
     auto out = cumheight + numerator / denominator;
-    output = where(clamp, input, out);
+    FVal<T> output_value = where(clamp, input, out);
 
     auto derivative_numerator = delta * delta *
         (derivative_plus_one * theta_theta + two_delta * theta_one_minus_theta +
          derivative * one_minus_theta * one_minus_theta);
-    det = where(clamp, 1., derivative_numerator / (denominator * denominator));
+    FVal<T> det_value =
+        where(clamp, 1., derivative_numerator / (denominator * denominator));
+    return {output_value, det_value};
 }
 
 template <typename T>
-KERNELSPEC void backward_kernel_rqs_forward(
+KERNELSPEC void kernel_rqs_forward(
+    FIn<T, 0> input, FIn<T, 1> condition, FOut<T, 0> output, FOut<T, 0> det
+) {
+    auto result = rqs_forward_values<T>(input, condition);
+    output = result.output;
+    det = result.det;
+}
+
+// the same, with the output also written in single precision
+template <typename T>
+KERNELSPEC void kernel_rqs_forward_with_float32(
     FIn<T, 0> input,
     FIn<T, 1> condition,
-    FIn<T, 0> output_grad,
-    FIn<T, 0> det_grad,
+    FOut<T, 0> output,
+    FOut<T, 0> det,
+    FOut<typename T::F32, 0> output32
+) {
+    auto result = rqs_forward_values<T>(input, condition);
+    output = result.output;
+    det = result.det;
+    output32 = FVal<typename T::F32>(result.output);
+}
+
+template <typename T>
+KERNELSPEC void backward_rqs_forward_impl(
+    FIn<T, 0> input,
+    FIn<T, 1> condition,
+    FVal<T> output_grad,
+    FVal<T> det_grad,
     FOut<T, 0> input_grad,
     FOut<T, 1> condition_grad
 ) {
@@ -386,9 +533,47 @@ KERNELSPEC void backward_kernel_rqs_forward(
 }
 
 template <typename T>
-KERNELSPEC void kernel_rqs_inverse(
-    FIn<T, 0> input, FIn<T, 1> condition, FOut<T, 0> output, FOut<T, 0> det
+KERNELSPEC void backward_kernel_rqs_forward(
+    FIn<T, 0> input,
+    FIn<T, 1> condition,
+    FIn<T, 0> output_grad,
+    FIn<T, 0> det_grad,
+    FOut<T, 0> input_grad,
+    FOut<T, 1> condition_grad
 ) {
+    backward_rqs_forward_impl<T>(
+        input,
+        condition,
+        FVal<T>(output_grad),
+        FVal<T>(det_grad),
+        input_grad,
+        condition_grad
+    );
+}
+
+// the gradient of the single-precision copy of the output adds to the output gradient
+template <typename T>
+KERNELSPEC void backward_kernel_rqs_forward_with_float32(
+    FIn<T, 0> input,
+    FIn<T, 1> condition,
+    FIn<T, 0> output_grad,
+    FIn<T, 0> det_grad,
+    FIn<typename T::F32, 0> output32_grad,
+    FOut<T, 0> input_grad,
+    FOut<T, 1> condition_grad
+) {
+    backward_rqs_forward_impl<T>(
+        input,
+        condition,
+        FVal<T>(output_grad) + FVal<T>(FVal<typename T::F32>(output32_grad)),
+        FVal<T>(det_grad),
+        input_grad,
+        condition_grad
+    );
+}
+
+template <typename T>
+KERNELSPEC RqsValues<T> rqs_inverse_values(FIn<T, 0> input, FIn<T, 1> condition) {
     FVal<T> height(condition[0]), width(condition[1]);
     FVal<T> cumheight(condition[2]), cumwidth(condition[3]);
     FVal<T> derivative_unorm(condition[4]), derivative_plus_one_unorm(condition[5]);
@@ -422,7 +607,7 @@ KERNELSPEC void kernel_rqs_inverse(
     auto discriminant = b * b + 4. * a * c;
     auto theta = 2. * c / (b + sqrt(discriminant));
     auto out = cumwidth + theta * width;
-    output = where(clamp, input, out);
+    FVal<T> output_value = where(clamp, input, out);
 
     auto one_minus_theta = 1. - theta;
     auto theta_one_minus_theta = theta * one_minus_theta;
@@ -430,15 +615,41 @@ KERNELSPEC void kernel_rqs_inverse(
     auto derivative_numerator = delta * delta *
         (derivative_plus_one * theta * theta + two_delta * theta_one_minus_theta +
          derivative * one_minus_theta * one_minus_theta);
-    det = where(clamp, 1., denominator * denominator / derivative_numerator);
+    FVal<T> det_value =
+        where(clamp, 1., denominator * denominator / derivative_numerator);
+    return {output_value, det_value};
 }
 
 template <typename T>
-KERNELSPEC void backward_kernel_rqs_inverse(
+KERNELSPEC void kernel_rqs_inverse(
+    FIn<T, 0> input, FIn<T, 1> condition, FOut<T, 0> output, FOut<T, 0> det
+) {
+    auto result = rqs_inverse_values<T>(input, condition);
+    output = result.output;
+    det = result.det;
+}
+
+// the same, with the output also written in single precision
+template <typename T>
+KERNELSPEC void kernel_rqs_inverse_with_float32(
     FIn<T, 0> input,
     FIn<T, 1> condition,
-    FIn<T, 0> output_grad,
-    FIn<T, 0> det_grad,
+    FOut<T, 0> output,
+    FOut<T, 0> det,
+    FOut<typename T::F32, 0> output32
+) {
+    auto result = rqs_inverse_values<T>(input, condition);
+    output = result.output;
+    det = result.det;
+    output32 = FVal<typename T::F32>(result.output);
+}
+
+template <typename T>
+KERNELSPEC void backward_rqs_inverse_impl(
+    FIn<T, 0> input,
+    FIn<T, 1> condition,
+    FVal<T> output_grad,
+    FVal<T> det_grad,
     FOut<T, 0> input_grad,
     FOut<T, 1> condition_grad
 ) {
@@ -600,6 +811,46 @@ KERNELSPEC void backward_kernel_rqs_inverse(
     condition_grad[5] = where(clamp, 0., grad_derivative_plus_one_unorm);
 }
 
+template <typename T>
+KERNELSPEC void backward_kernel_rqs_inverse(
+    FIn<T, 0> input,
+    FIn<T, 1> condition,
+    FIn<T, 0> output_grad,
+    FIn<T, 0> det_grad,
+    FOut<T, 0> input_grad,
+    FOut<T, 1> condition_grad
+) {
+    backward_rqs_inverse_impl<T>(
+        input,
+        condition,
+        FVal<T>(output_grad),
+        FVal<T>(det_grad),
+        input_grad,
+        condition_grad
+    );
+}
+
+// the gradient of the single-precision copy of the output adds to the output gradient
+template <typename T>
+KERNELSPEC void backward_kernel_rqs_inverse_with_float32(
+    FIn<T, 0> input,
+    FIn<T, 1> condition,
+    FIn<T, 0> output_grad,
+    FIn<T, 0> det_grad,
+    FIn<typename T::F32, 0> output32_grad,
+    FOut<T, 0> input_grad,
+    FOut<T, 1> condition_grad
+) {
+    backward_rqs_inverse_impl<T>(
+        input,
+        condition,
+        FVal<T>(output_grad) + FVal<T>(FVal<typename T::F32>(output32_grad)),
+        FVal<T>(det_grad),
+        input_grad,
+        condition_grad
+    );
+}
+
 template <typename T>
 KERNELSPEC void kernel_softmax(FIn<T, 1> input, FOut<T, 1> output) {
     FVal<T> norm(0.), in_max(0.);
diff --git a/madspace/src/kernels/operations.hpp b/madspace/src/kernels/operations.hpp
index b129ef6e8f..73eb2ef158 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,13 @@ void backward_batch_foreach(
     foreach_func(args, input_grads, batch_size, device);
 }
 
+// integers carry no gradient, so structural ops skip them (all their inputs share
+// one dtype)
+template <typename I>
+inline bool skip_int_grad(const I& instruction, const TensorVec& locals) {
+    return !is_float(locals[instruction.input_indices[0]].dtype());
+}
+
 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 +135,15 @@ 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)
+    if (skip_int_grad(instruction, locals)) {
+        return;
+    }
     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,
+                locals[input_index].dtype(),
                 locals[input_index].shape(),
                 device,
                 instruction.input_grad_alloc_hints[i]
@@ -182,11 +190,14 @@ template <typename I, typename D>
 void backward_op_unstack(
     const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device
 ) {
+    if (skip_int_grad(instruction, locals)) {
+        return;
+    }
     auto input_index = instruction.input_indices[0];
     auto& input_grad = local_grads[input_index];
     if (!input_grad) {
         input_grad = Tensor(
-            DataType::dt_float,
+            locals[input_index].dtype(),
             locals[input_index].shape(),
             device,
             instruction.input_grad_alloc_hints[0]
@@ -212,11 +223,14 @@ template <typename I, typename D>
 void backward_op_pop(
     const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device
 ) {
+    if (skip_int_grad(instruction, locals)) {
+        return;
+    }
     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,
+            locals[input_index].dtype(),
             locals[input_index].shape(),
             device,
             instruction.input_grad_alloc_hints[0]
@@ -264,12 +278,15 @@ template <typename I, typename D>
 void backward_op_batch_cat(
     const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device
 ) {
+    if (skip_int_grad(instruction, locals)) {
+        return;
+    }
     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,
+                locals[input_index].dtype(),
                 locals[input_index].shape(),
                 device,
                 instruction.input_grad_alloc_hints[i]
@@ -300,12 +317,15 @@ template <typename I, typename D>
 void backward_op_batch_split(
     const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device
 ) {
+    if (skip_int_grad(instruction, locals)) {
+        return;
+    }
     auto& sizes = locals[instruction.input_indices[1]].batch_sizes();
     auto input_index = instruction.input_indices[0];
     auto& input_grad = local_grads[input_index];
     if (!input_grad) {
         input_grad = Tensor(
-            DataType::dt_float,
+            locals[input_index].dtype(),
             locals[input_index].shape(),
             device,
             instruction.input_grad_alloc_hints[0]
@@ -352,12 +372,15 @@ template <typename I, typename D>
 void backward_op_cat(
     const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device
 ) {
+    if (skip_int_grad(instruction, locals)) {
+        return;
+    }
     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,
+                locals[input_index].dtype(),
                 locals[input_index].shape(),
                 device,
                 instruction.input_grad_alloc_hints[i]
@@ -465,13 +488,16 @@ template <typename I, typename D>
 void backward_op_rqs_reshape(
     const I& instruction, TensorVec& locals, TensorVec& local_grads, const D& device
 ) {
+    if (skip_int_grad(instruction, locals)) {
+        return;
+    }
     auto n_dims = instruction.output_shapes[0][0];
     auto n_bins = instruction.output_shapes[0][1];
     auto input_index = instruction.input_indices[0];
     auto& input_grad = local_grads[input_index];
     if (!input_grad) {
         input_grad = Tensor(
-            DataType::dt_float,
+            locals[input_index].dtype(),
             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 5313bb2a0b..ac91484e74 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,13 +78,44 @@ 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]);
     }
 }
 
+// select with a single-precision copy of the gathered entries
+template <typename T>
+KERNELSPEC void kernel_select_with_float32(
+    FIn<T, 1> input,
+    ISIn<T, 1> indices,
+    FOut<T, 1> output,
+    FOut<typename T::F32, 1> output32
+) {
+    for (std::size_t i = 0; i < indices.size(); ++i) {
+        auto value = input.gather(indices[i]);
+        output[i] = value;
+        output32[i] = FVal<typename T::F32>(value);
+    }
+}
+
+template <typename T>
+KERNELSPEC void backward_kernel_select_with_float32(
+    ISIn<T, 1> indices,
+    FIn<T, 1> output_grad,
+    FIn<typename T::F32, 1> output32_grad,
+    FOut<T, 1> input_grad,
+    IOut<T, 1> indices_grad
+) {
+    for (std::size_t i = 0; i < indices.size(); ++i) {
+        input_grad.scatter_add(
+            indices[i],
+            FVal<T>(output_grad[i]) + FVal<T>(FVal<typename T::F32>(output32_grad[i]))
+        );
+    }
+}
+
 template <typename T>
 KERNELSPEC void
 kernel_select_vector(FIn<T, 2> input, IIn<T, 1> indices, FOut<T, 2> output) {
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 4565c753bd..b868fcfb61 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",
@@ -173,9 +174,9 @@ Flow::Flow(
     // one-dimensional input (a 2 -> 1 phase space) would leave the flow
     // without any block and nothing to train.
     if (input_dim < 2) {
-        throw std::invalid_argument(std::format(
-            "Flow input dimension must be at least 2, got {}", input_dim
-        ));
+        throw std::invalid_argument(
+            std::format("Flow input dimension must be at least 2, got {}", input_dim)
+        );
     }
     std::size_t block_count = 0;
     for (std::size_t dim = input_dim - 1; dim > 0; dim /= 2) {
@@ -203,13 +204,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 5533fc48f1..4bef488e61 100644
--- a/madspace/src/phasespace/integrand.cpp
+++ b/madspace/src/phasespace/integrand.cpp
@@ -48,8 +48,8 @@ bool emits_alphas_scales(
     const std::optional<RunningCoupling>& running_coupling,
     std::size_t particle_count
 ) {
-    return energy_scale && energy_scale->mlm_alphas_reweighting() &&
-        running_coupling && particle_count > 3;
+    return energy_scale && energy_scale->mlm_alphas_reweighting() && running_coupling &&
+        particle_count > 3;
 }
 
 // Whether the event carries the MLM pdf reweighting ladder, for the
@@ -212,8 +212,7 @@ Integrand::Integrand(
                         energy_scale, running_coupling, mapping.particle_count()
                     )) {
                     ret_types.push_back(
-                        "alphas_scales",
-                        batch_float_array(mapping.particle_count() - 3)
+                        "alphas_scales", batch_float_array(mapping.particle_count() - 3)
                     );
                 }
                 if (emits_pdf_ladder(
@@ -374,19 +373,15 @@ Integrand::Integrand(
         for (std::size_t f = 0; f < option_count; ++f) {
             _pdf_rw_table.at(0 * option_count + f) = index_of(21);
             for (std::size_t beam = 0; beam < 2; ++beam) {
-                _pdf_rw_table.at((1 + beam) * option_count + f) = index_of(
-                    pid_options.size() > 0 ? pid_options.at(f).at(beam) : 21
-                );
+                _pdf_rw_table.at((1 + beam) * option_count + f) =
+                    index_of(pid_options.size() > 0 ? pid_options.at(f).at(beam) : 21);
             }
             for (std::size_t c = 0; c < absolute.size(); ++c) {
-                _pdf_rw_table.at((3 + c) * option_count + f) =
-                    index_of(absolute.at(c));
+                _pdf_rw_table.at((3 + c) * option_count + f) = index_of(absolute.at(c));
             }
         }
         for (std::size_t c = 0; c < class_count; ++c) {
-            _pdf_rw_class_offsets.push_back(
-                static_cast<me_int_t>(c * option_count)
-            );
+            _pdf_rw_class_offsets.push_back(static_cast<me_int_t>(c * option_count));
         }
         _pdf_rw = PartonDensity(*grid, rw_pids, true);
         _pdf_rw_pdgs.assign(rw_pids.begin(), rw_pids.end());
@@ -492,13 +487,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;
@@ -522,7 +517,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);
@@ -553,7 +548,9 @@ NamedVector<Type> Integrand::compute_channel_part_ret_types() const {
     }
     if (_pdf_rw && _pid_options.size() > 0) {
         for (auto name : ladder_slot_names) {
-            ret.push_back(std::format("{}_acc", name), acc_float_array(particle_count - 2));
+            ret.push_back(
+                std::format("{}_acc", name), acc_float_array(particle_count - 2)
+            );
         }
         for (auto name : ladder_beam_names) {
             ret.push_back(std::format("{}_acc", name), acc_float_array(2));
@@ -727,10 +724,11 @@ NamedVector<Value> Integrand::build_channel_part(
             _pdfs.at(i).value().build_function(fb, {x_acc.at(i), pdf_scale}).at(0);
         if (_has_mirror && _pdfs_swapped.at(i)) {
             // leg i of a mirrored event comes from the other beam
-            auto pdf_swapped = _pdfs_swapped.at(i)
-                                   .value()
-                                   .build_function(fb, {x_acc.at(i), pdf_scale})
-                                   .at(0);
+            auto pdf_swapped =
+                _pdfs_swapped.at(i)
+                    .value()
+                    .build_function(fb, {x_acc.at(i), pdf_scale})
+                    .at(0);
             pdf = fb.gather_vector(
                 fb.batch_gather(indices_acc, mirror_index), fb.stack({pdf, pdf_swapped})
             );
@@ -956,8 +954,9 @@ NamedVector<Value> Integrand::build_channel_part(
         for (auto name : ladder_slot_names) {
             out.push_back(
                 std::format("{}_acc", name),
-                mlm_ladder.size() > 0 ? mlm_ladder.at(name)
-                                      : fb.full({0., batch_size, slot_count})
+                mlm_ladder.size() > 0
+                    ? mlm_ladder.at(name)
+                    : fb.full({0., batch_size, slot_count})
             );
         }
         for (auto name : ladder_beam_names) {
@@ -1034,9 +1033,8 @@ ValueVec Integrand::mlm_weights(
         auto flavor_classes = scales.at("pdf_rw_flavor");
         std::size_t slot_count = flavor_classes.type.shape.at(0);
         double low = _energy_scale->min_scale();
-        double high = _energy_scale->max_scale() > 0.
-            ? _energy_scale->max_scale()
-            : 1e30;
+        double high =
+            _energy_scale->max_scale() > 0. ? _energy_scale->max_scale() : 1e30;
         auto ones = fb.full({1., batch_size_acc});
         Value factor, veto;
         ValueVec ladder_x, ladder_pdg, ladder_q, ladder_info;
@@ -1052,10 +1050,7 @@ ValueVec Integrand::mlm_weights(
             auto beam = pick("pdf_rw_beam");
             // x of the beam line: the beam's own momentum fraction times
             // everything the clustering z has taken off it since.
-            auto x_beam = fb.add(
-                x.at(0),
-                fb.mul(beam, fb.sub(x.at(1), x.at(0)))
-            );
+            auto x_beam = fb.add(x.at(0), fb.mul(beam, fb.sub(x.at(1), x.at(0))));
             auto x = fb.mul(x_beam, pick("pdf_rw_x"));
             auto scale_of = [&](const char* name) {
                 auto q = pick(name);
@@ -1123,7 +1118,8 @@ ValueVec Integrand::mlm_weights(
             // each beam's leading density, at the scale it was evaluated at
             // (clamped by apply_scale_range)
             ladder->push_back(
-                "ladder_lead_q", fb.stack({scales.at("pdf_scale1"), scales.at("pdf_scale2")})
+                "ladder_lead_q",
+                fb.stack({scales.at("pdf_scale1"), scales.at("pdf_scale2")})
             );
         }
     }
@@ -1239,8 +1235,9 @@ NamedVector<Value> Integrand::build_common_part(
     if (_energy_scale && _energy_scale->is_mlm()) {
         cluster_scales_acc = args.at("cluster_scales_acc");
     }
-    bool has_alphas_scales =
-        emits_alphas_scales(_energy_scale, _running_coupling, _mapping.particle_count());
+    bool has_alphas_scales = emits_alphas_scales(
+        _energy_scale, _running_coupling, _mapping.particle_count()
+    );
     Value alphas_scales_acc;
     if (has_alphas_scales) {
         alphas_scales_acc = args.at("alphas_scales_acc");
@@ -1409,20 +1406,19 @@ NamedVector<Value> Integrand::build_common_part(
                 if (_diff_xs.at(0).has_pdf(i)) {
                     // a first density of exactly zero gave a zero weight, which
                     // stays zero (the guard only keeps 0 * x / 0 from a NaN)
-                    auto denominator = fb.add(
-                        pdfs_first.at(i), fb.cut_one(pdfs_first.at(i), 0., 0.)
-                    );
-                    factor =
-                        fb.mul(factor, fb.div(pdfs_acc.at(i), denominator));
+                    auto denominator =
+                        fb.add(pdfs_first.at(i), fb.cut_one(pdfs_first.at(i), 0., 0.));
+                    factor = fb.mul(factor, fb.div(pdfs_acc.at(i), denominator));
                 }
             }
             dxs_vec.at(0) = fb.mul(dxs_vec.at(0), factor);
         } else {
             // several powers of alpha_s mix (gluon fusion and VBF in one
             // process, say): the matrix element has to be evaluated again
-            dxs_vec = _diff_xs.at(0)
-                          .build_function(fb, make_xs_args(pdfs_acc, alpha_qcd_acc))
-                          .values();
+            dxs_vec =
+                _diff_xs.at(0)
+                    .build_function(fb, make_xs_args(pdfs_acc, alpha_qcd_acc))
+                    .values();
         }
     }
 
@@ -1605,9 +1601,7 @@ NamedVector<Value> Integrand::build_common_part(
             );
         }
 
-        outputs.push_back(
-            "ren_scale", scatter_or_drop(zeros_float, ren_scale_acc)
-        );
+        outputs.push_back("ren_scale", scatter_or_drop(zeros_float, ren_scale_acc));
         outputs.push_back("alpha_qcd", scatter_or_drop(zeros_float, alpha_qcd_acc));
         ValueVec pdf_vals;
         if (_diff_xs.at(0).has_pdf(0)) {
@@ -1636,8 +1630,7 @@ NamedVector<Value> Integrand::build_common_part(
             outputs.push_back(
                 "cluster_scales",
                 scatter_or_drop(
-                    fb.full({0., batch_size_val, outgoing_count}),
-                    cluster_scales_acc
+                    fb.full({0., batch_size_val, outgoing_count}), cluster_scales_acc
                 )
             );
         }
@@ -1646,8 +1639,7 @@ NamedVector<Value> Integrand::build_common_part(
                 auto& value = ladder_acc.at(name);
                 auto width = static_cast<me_int_t>(value.type.shape.at(0));
                 outputs.push_back(
-                    name,
-                    scatter_or_drop(fb.full({0., batch_size_val, width}), value)
+                    name, scatter_or_drop(fb.full({0., batch_size_val, width}), value)
                 );
             }
         }
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 fc3a28025c..944bbc3ebe 100644
--- a/madspace/src/phasespace/pdf.cpp
+++ b/madspace/src/phasespace/pdf.cpp
@@ -322,14 +322,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();
@@ -339,9 +339,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);
@@ -442,17 +442,17 @@ NamedVector<Value> PartonDensity::build_function_impl(
     }
     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) {
@@ -597,11 +597,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();
@@ -610,8 +610,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);
@@ -635,12 +635,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 0a7e43a66f..c50c0349aa 100644
--- a/madspace/src/phasespace/phasespace.cpp
+++ b/madspace/src/phasespace/phasespace.cpp
@@ -257,7 +257,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 6f233039d5..e5d1425cf8 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 9a46c0c1a1..8c761ecd3d 100644
--- a/madspace/src/python/madspace.cpp
+++ b/madspace/src/python/madspace.cpp
@@ -7,6 +7,7 @@
 #include "function_runtime.hpp"
 #include "instruction_set.hpp"
 #include "madspace/compgraphs.hpp"
+#include "madspace/compgraphs/optimizer.hpp"
 #include "madspace/driver.hpp"
 #include "madspace/phasespace.hpp"
 #include "madspace/util.hpp"
@@ -149,8 +150,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},
         },
         "",
@@ -284,6 +286,15 @@ PYBIND11_MODULE(_madspace_py, m) {
             pydoc::doc("Function::instructions")
         );
 
+    m.def(
+        "fuse_instructions",
+        &fuse_instructions,
+        py::arg("function"),
+        "Fuses single-precision conversions and network activations into the "
+        "instructions producing their inputs (the lowering the GPU backend applies "
+        "when it builds a runtime)."
+    );
+
     py::classh<Device> device(m, "Device", pydoc::doc("Device"));
     m.def(
         "cpu_device",
@@ -1669,17 +1680,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,
@@ -1704,7 +1718,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") = "",
@@ -1713,6 +1728,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"))
@@ -1791,7 +1807,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,
@@ -1799,6 +1816,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"))
@@ -1868,7 +1886,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>{},
@@ -1876,6 +1895,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(
@@ -1947,57 +1967,50 @@ PYBIND11_MODULE(_madspace_py, m) {
             {"goodjet", PartonLineScheme::goodjet},
         }
     );
-    mlm
-        .def(
-            py::init<
-                std::vector<Topology>,
-                nested_vector3<std::size_t>,
-                nested_vector2<std::size_t>,
-                double,
-                JetScaleScheme,
-                ScaleScheme,
-                std::unordered_map<int, int>,
-                double,
-                double,
-                double,
-                bool,
-                std::vector<int>,
-                int,
-                PartonLineScheme,
-                AlphasScheme,
-                bool,
-                ClusteringMeasure,
-                ClusteringHistory,
-                bool>(),
-            py::arg("topologies"),
-            py::arg("permutations"),
-            py::arg("diagram_indices"),
-            py::arg("cm_energy"),
-            py::arg("jet_scale_scheme") = JetScaleScheme::production,
-            py::arg("scale_scheme") = ScaleScheme::clustering_mean,
-            py::arg("pdg_color_types") = std::unordered_map<int, int>{},
-            py::arg("xqcut") = 0.,
-            py::arg("bw_cutoff") = 15.,
-            py::arg("jet_radius") = 1.0,
-            py::arg("hadronic") = true,
-            py::arg("external_pdg_ids") = std::vector<int>{},
-            py::arg("max_jet_flavor") = 4,
-            py::arg("parton_line_scheme") = PartonLineScheme::goodjet,
-            py::arg("alphas_scheme") = AlphasScheme::per_vertex,
-            py::arg("pdf_reweighting") = true,
-            py::arg("clustering_measure") = ClusteringMeasure::fxfx,
-            py::arg("clustering_history") = ClusteringHistory::all_diagrams,
-            py::arg("clustering_mean_pdf_ladder") = true
-        )
-        .def_property_readonly(
-            "pdf_absolute_pdgs", &MLMClustering::pdf_absolute_pdgs
-        )
-        .def_property_readonly(
-            "clustering_measure", &MLMClustering::clustering_measure
-        )
-        .def_property_readonly(
-            "clustering_history", &MLMClustering::clustering_history
-        )
+    mlm.def(
+           py::init<
+               std::vector<Topology>,
+               nested_vector3<std::size_t>,
+               nested_vector2<std::size_t>,
+               double,
+               JetScaleScheme,
+               ScaleScheme,
+               std::unordered_map<int, int>,
+               double,
+               double,
+               double,
+               bool,
+               std::vector<int>,
+               int,
+               PartonLineScheme,
+               AlphasScheme,
+               bool,
+               ClusteringMeasure,
+               ClusteringHistory,
+               bool>(),
+           py::arg("topologies"),
+           py::arg("permutations"),
+           py::arg("diagram_indices"),
+           py::arg("cm_energy"),
+           py::arg("jet_scale_scheme") = JetScaleScheme::production,
+           py::arg("scale_scheme") = ScaleScheme::clustering_mean,
+           py::arg("pdg_color_types") = std::unordered_map<int, int>{},
+           py::arg("xqcut") = 0.,
+           py::arg("bw_cutoff") = 15.,
+           py::arg("jet_radius") = 1.0,
+           py::arg("hadronic") = true,
+           py::arg("external_pdg_ids") = std::vector<int>{},
+           py::arg("max_jet_flavor") = 4,
+           py::arg("parton_line_scheme") = PartonLineScheme::goodjet,
+           py::arg("alphas_scheme") = AlphasScheme::per_vertex,
+           py::arg("pdf_reweighting") = true,
+           py::arg("clustering_measure") = ClusteringMeasure::fxfx,
+           py::arg("clustering_history") = ClusteringHistory::all_diagrams,
+           py::arg("clustering_mean_pdf_ladder") = true
+    )
+        .def_property_readonly("pdf_absolute_pdgs", &MLMClustering::pdf_absolute_pdgs)
+        .def_property_readonly("clustering_measure", &MLMClustering::clustering_measure)
+        .def_property_readonly("clustering_history", &MLMClustering::clustering_history)
         .def_property_readonly(
             "diagram_start_states", &MLMClustering::diagram_start_states
         )
@@ -3715,9 +3728,7 @@ PYBIND11_MODULE(_madspace_py, m) {
             "is_scale", &Variation::is_scale, pydoc::doc("Variation::is_scale")
         )
         .def_property_readonly(
-            "is_emission",
-            &Variation::is_emission,
-            pydoc::doc("Variation::is_emission")
+            "is_emission", &Variation::is_emission, pydoc::doc("Variation::is_emission")
         );
     py::classh<PdfGroupInfo>(m, "PdfGroupInfo", pydoc::doc("PdfGroupInfo"))
         .def_readonly(
@@ -3862,10 +3873,12 @@ PYBIND11_MODULE(_madspace_py, m) {
                 int particle_flags =
                     (momenta.empty() ? ParticleRecord::f_none
                                      : ParticleRecord::f_particle_data) |
-                    (alphas_scales.empty() ? ParticleRecord::f_none
-                                           : ParticleRecord::f_alphas_scales) |
-                    (ladder_lead.empty() ? ParticleRecord::f_none
-                                         : ParticleRecord::f_pdf_ladder);
+                    (alphas_scales.empty()
+                         ? ParticleRecord::f_none
+                         : ParticleRecord::f_alphas_scales) |
+                    (ladder_lead.empty()
+                         ? ParticleRecord::f_none
+                         : ParticleRecord::f_pdf_ladder);
                 DataLayout layout(
                     EventRecord::layout(
                         EventRecord::f_weight | EventRecord::f_subproc_index |
@@ -3900,7 +3913,8 @@ PYBIND11_MODULE(_madspace_py, m) {
                         particle.py() = p.at(2);
                         particle.pz() = p.at(3);
                     }
-                    for (std::size_t j = 0; !alphas_scales.empty() && j < particle_count;
+                    for (std::size_t j = 0;
+                         !alphas_scales.empty() && j < particle_count;
                          ++j) {
                         auto& row = alphas_scales.at(i);
                         buffer.particle(i, j).alphas_scale() =
@@ -3912,8 +3926,8 @@ PYBIND11_MODULE(_madspace_py, m) {
                         auto& lead = ladder_lead.at(i);
                         particle.ladder_lead_q() = j < 2 ? lead.at(j) : 0.;
                         bool has_step = !ladder.empty() && j < ladder.at(i).size();
-                        auto step = has_step ? ladder.at(i).at(j)
-                                             : std::vector<double>(4, 0.);
+                        auto step =
+                            has_step ? ladder.at(i).at(j) : std::vector<double>(4, 0.);
                         particle.ladder_x() = step.at(0);
                         particle.ladder_pdg() = step.at(1);
                         particle.ladder_q() = step.at(2);
diff --git a/madspace/src/python/python_function.cpp b/madspace/src/python/python_function.cpp
index 4f044ad089..f9d8f51e3e 100644
--- a/madspace/src/python/python_function.cpp
+++ b/madspace/src/python/python_function.cpp
@@ -63,7 +63,7 @@ Tensor tensor_view(
         throw std::invalid_argument(std::format("unknown device {}", device_name));
     }
     std::size_t item_size =
-        dtype == DataType::dt_float ? sizeof(double) : sizeof(me_int_t);
+        dtype == DataType::dt_float64 ? sizeof(double) : sizeof(me_int_t);
     void* data = reinterpret_cast<void*>(address + offset * item_size);
     return Tensor(dtype, Sizes(shape), Sizes(stride), device, data, [] {});
 }
diff --git a/madspace/tests/test_float32.py b/madspace/tests/test_float32.py
new file mode 100644
index 0000000000..88b9f9aae1
--- /dev/null
+++ b/madspace/tests/test_float32.py
@@ -0,0 +1,239 @@
+"""Single-precision NN instructions: forward and backward against torch."""
+
+import pytest
+import torch
+from madspace_testing import DEVICE, default_context, new_context
+
+import madspace as ms
+from madspace.torch import FunctionModule
+
+
+@pytest.fixture
+def gen():
+    """A fresh local generator per test, so results don't depend on test order."""
+    return torch.Generator().manual_seed(1234)
+
+
+def _draw(draw, *args, requires_grad=False, **kwargs):
+    """Random numbers from the CPU generator, moved to the device of the tests."""
+    values = draw(*args, device="cpu", **kwargs).to(DEVICE)
+    return values.requires_grad_(requires_grad)
+
+
+def _randn(*args, **kwargs):
+    return _draw(torch.randn, *args, **kwargs)
+
+
+def _rand(*args, **kwargs):
+    return _draw(torch.rand, *args, **kwargs)
+
+
+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, gen):
+    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(), default_context())
+
+    x = (3 * _randn((N, DIMS_IN), dtype=torch_dtype, generator=gen)).requires_grad_()
+    out = module(x)
+    assert out.dtype == torch_dtype
+    out_grad = _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, gen):
+    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(), default_context())
+
+    x = _randn((N, DIMS_IN), dtype=torch_dtype, generator=gen, requires_grad=True)
+    # single inputs are passed without batch dimension
+    w = _randn(
+        (DIMS_OUT, DIMS_IN), dtype=torch_dtype, generator=gen, requires_grad=True
+    )
+    b = _randn((DIMS_OUT,), dtype=torch_dtype, generator=gen, requires_grad=True)
+    y = module(x, w, b)
+    assert y.dtype == torch_dtype
+    y_grad = _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(gen):
+    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(), default_context())
+
+    x = _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(gen):
+    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(), default_context())
+    x = _randn((N, 1), dtype=torch.float32, generator=gen)
+    torch.testing.assert_close(module(x), 2 * x + 1)
+
+
+def test_flow_float32_subnets(tmp_path, gen):
+    """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 = new_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 * _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 = _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-3, atol=1e-4)
+    for g32, g64 in zip(grads[ms.DataType.float32], grads[ms.DataType.float64]):
+        torch.testing.assert_close(g32, g64, rtol=5e-3, atol=1e-4)
+
+    # float32 globals survive saving and loading
+    ctx.save_globals(str(tmp_path / "globals.tar"))
+    loaded = new_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(gen):
+    """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(), default_context())
+
+    x = _randn((N, 3), dtype=torch.float32, generator=gen, requires_grad=True)
+    out = module(x)
+    out_grad = _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)
+
+
+def test_random_int_gather(gen):
+    """random_int returns int32 indices, as used when sampling the replay buffer."""
+    io_type = batch_type(ms.DataType.float64, [])
+    fb = ms.FunctionBuilder(
+        ms.NamedTypes([("in", io_type)]), ms.NamedTypes([("out", io_type)])
+    )
+    size = fb.batch_size([fb.input(0)])
+    fb.output(0, fb.batch_gather(fb.random_int(size, size), fb.input(0)))
+    module = FunctionModule(fb.function(), default_context())
+    x = _rand((N,), dtype=torch.float64, generator=gen)
+    out = module(x)
+    assert out.shape == x.shape
+    assert torch.isin(out, x).all()
diff --git a/madspace/tests/test_fused_float32.py b/madspace/tests/test_fused_float32.py
new file mode 100644
index 0000000000..4480038301
--- /dev/null
+++ b/madspace/tests/test_fused_float32.py
@@ -0,0 +1,222 @@
+"""Fused single-precision instructions against their unfused equivalents."""
+
+import re
+
+import pytest
+import torch
+from madspace_testing import ON_GPU, generator, new_context, rand, randn
+
+import madspace as ms
+from madspace.torch import FunctionModule
+
+N = 211
+DIMS_IN = 7
+DIMS_OUT = 5
+
+
+def batch_type(dtype, shape):
+    return ms.Type(dtype, ms.batch_size, shape)
+
+
+def matmul_module(fused):
+    f32, f64 = ms.DataType.float32, ms.DataType.float64
+    fb = ms.FunctionBuilder(
+        ms.NamedTypes(
+            [
+                ("x", batch_type(f32, [DIMS_IN])),
+                ("w", ms.Type(f32, ms.BatchSize.one, [DIMS_OUT, DIMS_IN])),
+                ("b", ms.Type(f32, ms.BatchSize.one, [DIMS_OUT])),
+            ]
+        ),
+        ms.NamedTypes([("y", batch_type(f64, [DIMS_OUT]))]),
+    )
+    x, w, b = fb.input(0), fb.input(1), fb.input(2)
+    if fused:
+        y = fb.matmul_to_float64(x, w, b)
+    else:
+        y = fb.convert_float64(fb.matmul(x, w, b))
+    fb.output(0, y)
+    return FunctionModule(fb.function(), new_context())
+
+
+def test_matmul_to_float64(monkeypatch):
+    g = generator(77)
+    x = randn((N, DIMS_IN), generator=g, dtype=torch.float32, requires_grad=True)
+    w = randn((DIMS_OUT, DIMS_IN), generator=g, dtype=torch.float32, requires_grad=True)
+    b = randn((DIMS_OUT,), generator=g, dtype=torch.float32, requires_grad=True)
+    y_grad = randn((N, DIMS_OUT), generator=g, dtype=torch.float64)
+
+    results = []
+    for fused in (True, False):
+        with monkeypatch.context() as mp:
+            # the GPU runtime would lower the unfused graph as well
+            if not fused:
+                mp.setenv("MADSPACE_NO_FUSION", "1")
+            xs, ws, bs = (t.detach().clone().requires_grad_() for t in (x, w, b))
+            y = matmul_module(fused)(xs, ws, bs)
+            assert y.dtype == torch.float64
+            (y * y_grad).sum().backward()
+            results.append((y.detach(), xs.grad, ws.grad, bs.grad))
+
+    y_ref = (x @ w.T + b).double()
+    torch.testing.assert_close(results[0][0], y_ref, rtol=1e-5, atol=1e-5)
+    # the fused output rounds exactly like matmul + convert_float64
+    torch.testing.assert_close(results[0][0], results[1][0], rtol=0, atol=0)
+    # the gradients through the fused and the unfused graph agree
+    for fused_t, unfused_t in zip(results[0], results[1]):
+        torch.testing.assert_close(fused_t, unfused_t, rtol=1e-4, atol=1e-4)
+
+    g_ref = y_grad.float()
+    torch.testing.assert_close(results[0][1], g_ref @ w.detach(), rtol=1e-4, atol=1e-4)
+    torch.testing.assert_close(
+        results[0][2], g_ref.T @ x.detach(), rtol=1e-4, atol=1e-3
+    )
+    torch.testing.assert_close(results[0][3], g_ref.sum(dim=0), rtol=1e-4, atol=1e-3)
+
+
+FUSED_NAMES = (
+    "convert_float32",
+    "convert_float64",
+    "matmul_float32",
+    "matmul_to_float64",
+    "leaky_relu_float32",
+    "matmul_leaky_relu_float32",
+    "select_with_float32",
+    "rqs_forward_with_float32",
+    "rqs_inverse_with_float32",
+)
+
+
+def instruction_counts(function):
+    names = "|".join(FUSED_NAMES)
+    counts = {}
+    for name in re.findall(rf"\b({names})\b", str(function)):
+        counts[name] = counts.get(name, 0) + 1
+    return counts
+
+
+def make_flow(prefix):
+    return ms.Flow(
+        4,
+        prefix=prefix,
+        subnet_precision=ms.DataType.float32,
+        subnet_hidden_dim=16,
+        subnet_layers=3,
+        invert_spline=False,
+    )
+
+
+@pytest.mark.parametrize("direction", ["forward", "inverse"])
+def test_conversion_lowering_graph(direction):
+    flow = make_flow("lowering")
+    function = getattr(flow, f"{direction}_function")()
+    before = instruction_counts(function)
+    after = instruction_counts(ms.fuse_instructions(function))
+    spline = "rqs_forward" if direction == "forward" else "rqs_inverse"
+    # every MLP of the flow converts its input and its output
+    assert before["convert_float32"] == before["convert_float64"] == 4
+    assert "convert_float32" not in after and "convert_float64" not in after
+    assert after["matmul_to_float64"] == 4
+    assert after["select_with_float32"] + after[f"{spline}_with_float32"] == 4
+    # the hidden layers fuse their activation
+    assert before["leaky_relu_float32"] == after["matmul_leaky_relu_float32"] == 8
+    assert "matmul_float32" not in after and "leaky_relu_float32" not in after
+    # the original function is not modified, and the pass is idempotent
+    assert instruction_counts(function) == before
+    assert (
+        instruction_counts(ms.fuse_instructions(ms.fuse_instructions(function)))
+        == after
+    )
+
+
+@pytest.mark.parametrize("direction", ["forward", "inverse"])
+def test_conversion_lowering_numerics(direction, monkeypatch):
+    """The lowered flow agrees with the unlowered one in values and gradients. On a
+    GPU the reference runs with the lowering disabled."""
+    g = generator(2024)
+    flow = make_flow(f"lowering_{direction}")
+    ctx = new_context()
+    flow.initialize_globals(ctx, 42)
+    # randomize all weights (the last layers start at zero)
+    for name in ctx.global_names():
+        weights = ctx.get_global(name).torch()
+        weights[:] = 0.3 * randn(weights.shape, generator=g, dtype=torch.float32).to(
+            weights.dtype
+        )
+
+    r = rand((137, 4), generator=g, dtype=torch.float64)
+    function = getattr(flow, f"{direction}_function")()
+    results = {}
+    for lowered in (True, False):
+        with monkeypatch.context() as mp:
+            if not lowered:
+                mp.setenv("MADSPACE_NO_FUSION", "1")
+            module = FunctionModule(
+                ms.fuse_instructions(function) if lowered else function, ctx
+            )
+            r_grad = r.clone().requires_grad_()
+            out, det = module(r_grad)
+            (out.sum() + det.log().sum()).backward()
+            results[lowered] = (
+                out.detach(),
+                det.detach(),
+                r_grad.grad,
+                {n: p.grad.clone() for n, p in module.global_params.items()},
+            )
+    for lowered_t, ref_t in zip(results[True][:3], results[False][:3]):
+        torch.testing.assert_close(lowered_t, ref_t, rtol=1e-12, atol=1e-12)
+    for name, grad in results[True][3].items():
+        torch.testing.assert_close(grad, results[False][3][name], rtol=1e-4, atol=1e-6)
+
+
+ACTIVATIONS = ("relu", "leaky_relu", "elu", "gelu", "sigmoid", "softplus")
+
+
+@pytest.mark.parametrize("precision", ["float32", "float64"])
+@pytest.mark.parametrize("activation", ACTIVATIONS)
+def test_matmul_activation(activation, precision, monkeypatch):
+    """An MLP with fused activations agrees with the unfused one. The two round
+    alike, except that the CPU runs the fused activation with the scalar instead of
+    the SIMD functions, which differ for the transcendental ones."""
+    g = generator(31)
+    dtype = getattr(ms.DataType, precision)
+    torch_dtype = getattr(torch, precision)
+    mlp = ms.MLP(
+        5, 3, 8, 3, getattr(ms.MLP, activation), f"act_{activation}_{precision}", dtype
+    )
+    ctx = new_context()
+    mlp.initialize_globals(ctx, 7)
+    for name in ctx.global_names():
+        weights = ctx.get_global(name).torch()
+        weights[:] = randn(weights.shape, generator=g, dtype=torch.float64).to(
+            weights.dtype
+        )
+    function = mlp.function()
+    lowered = ms.fuse_instructions(function)
+    assert instruction_counts_any(lowered, f"matmul_{activation}_{precision}") == 2
+    assert instruction_counts_any(lowered, f"{activation}_{precision}") == 0
+
+    x = randn((97, 5), generator=g, dtype=torch.float64)
+    y_grad = randn((97, 3), generator=g, dtype=torch.float64)
+    results = {}
+    for fused in (True, False):
+        with monkeypatch.context() as mp:
+            if not fused:
+                mp.setenv("MADSPACE_NO_FUSION", "1")
+            module = FunctionModule(lowered if fused else function, ctx)
+            xs = x.clone().requires_grad_()
+            y = module(xs)
+            (y * y_grad).sum().backward()
+            results[fused] = [
+                y.detach(),
+                xs.grad,
+                *(p.grad.clone() for p in module.global_params.values()),
+            ]
+    exact = ON_GPU or activation in ("relu", "leaky_relu")
+    tol = 0 if exact else (1e-5 if torch_dtype == torch.float32 else 1e-12)
+    for fused_t, ref_t in zip(results[True], results[False]):
+        torch.testing.assert_close(fused_t, ref_t, rtol=tol, atol=tol)
+
+
+def instruction_counts_any(function, name):
+    return len(re.findall(rf"\b{name}\(", str(function)))
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..289fe28222 100644
--- a/madspace/tests/test_mlp.py
+++ b/madspace/tests/test_mlp.py
@@ -14,9 +14,21 @@ 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):
+    # smooth activation: at a kink, float32 rounding can flip the gradient of an event
+    return ms.MLP(10, 1, 32, 3, ms.MLP.gelu, "", precision[0])
 
 
 @pytest.fixture
@@ -34,9 +46,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 +92,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.GELU(),
         nn.Linear(32, 32),
-        nn.LeakyReLU(),
+        nn.GELU(),
         nn.Linear(32, 1),
-    )
+    ).to(dtype)
 
     with torch.no_grad():
         for i in range(3):
@@ -102,47 +123,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..045dd066b3 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 = "float32" # 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 d03e2e7ef7..4e900fb0cf 100755
--- a/tests/unit_tests/various/test_banner.py
+++ b/tests/unit_tests/various/test_banner.py
@@ -2378,6 +2378,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'], 'float32')
         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
@@ -2406,17 +2407,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'], 'float32')
 
     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'] = 'float64'
         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'], 'float64')
         self.assertEqual(rc2['multiparticles']['lepton'], rc['multiparticles']['lepton'])
         self.assertEqual(rc2['cuts']['jet-delta_r'], {'min': 0.4})