diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py
index 62626f3274..4a4a01f3ab 100644
--- a/madgraph/iolibs/template_files/mg7/gridpack.py
+++ b/madgraph/iolibs/template_files/mg7/gridpack.py
@@ -171,7 +171,7 @@ def main() -> None:
     event_generator = ms.EventGenerator(
         contexts=contexts,
         channels=channel_generators,
-        status_file=os.path.join(run_path, "info.json"),
+        status_file=ms.StatusFile(os.path.join(run_path, "info.json")),
         config=config,
     )
 
diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py
index 6114a4871b..f8622eba31 100644
--- a/madgraph/iolibs/template_files/mg7/madevent.py
+++ b/madgraph/iolibs/template_files/mg7/madevent.py
@@ -205,6 +205,7 @@ def init_event_dir(self) -> None:
                 break
             except FileExistsError:
                 run_index += 1
+        self.status_file = ms.StatusFile(os.path.join(self.run_path, "info.json"))
 
     def init_context(self) -> None:
         device_names = self.run_card["run"]["devices"]
@@ -424,7 +425,7 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera
         event_generator = ms.EventGenerator(
             contexts=self.contexts,
             channels=channel_generators,
-            status_file=os.path.join(self.run_path, "info.json"),
+            status_file=self.status_file,
             config=self.event_generator_config,
         )
         unused_globals = (
@@ -502,9 +503,6 @@ def train_madnis(self) -> None:
         madnis_args = self.run_card["madnis"]
         if not madnis_args["enable"]:
             return
-        if madnis_args.get("old", False):
-            self.train_madnis_old()
-            return
 
         gen_args = self.run_card["generation"]
         run_args = self.run_card["run"]
@@ -528,6 +526,7 @@ def train_madnis(self) -> None:
         config.adam_beta1 = madnis_args["adam_beta1"]
         config.adam_beta2 = madnis_args["adam_beta2"]
         config.adam_eps = madnis_args["adam_eps"]
+        config.grad_clip_threshold = madnis_args["grad_clip_threshold"]
         config.buffer_capacity = madnis_args["buffer_capacity"]
         config.minimum_buffer_size = madnis_args["minimum_buffer_size"]
         config.buffered_steps = madnis_args["buffered_steps"]
@@ -559,6 +558,7 @@ def train_madnis(self) -> None:
             config=config,
             integrands=integrands,
             cwnets=cwnets,
+            status_file=self.status_file,
         )
         madnis_training.train()
         for phasespace, active_channels in zip(
@@ -572,49 +572,6 @@ def train_madnis(self) -> None:
             context.copy_globals_from(self.contexts[0])
         self.event_generator = self.build_event_generator(madnis_phasespaces)
 
-    def train_madnis_old(self) -> None:
-        madnis_args = self.run_card["madnis"]
-        if not madnis_args["enable"]:
-            return
-
-        if len(self.subprocesses) > 1:
-            self.madnis_lower_box = ms.PrettyBox(
-                "Subprocesses", len(self.subprocesses) + 1, [12, 12, 12, 0],
-            )
-            self.madnis_lower_box.set_row(0, ["Subprocess", "Loss", "Channels", "Batch"])
-            self.madnis_upper_box = ms.PrettyBox(
-                "MadNIS training", 2, [18, 0], self.madnis_lower_box.line_count
-            )
-            self.madnis_upper_box.set_column(0, ["Subprocesses:", "Run time:"])
-            self.madnis_upper_box.print_first()
-            self.madnis_lower_box.print_first()
-        else:
-            self.madnis_box = ms.PrettyBox(
-                "MadNIS training", 4, [18, 0]
-            )
-            self.madnis_box.set_column(0, ["Batch:", "Loss:", "Channels:", "Run time:"])
-            self.madnis_box.print_first()
-
-        self.last_update_time = 0
-        self.madnis_wall_time = time.time()
-        self.madnis_cpu_time = time.process_time()
-
-        madnis_phasespaces = []
-        for subproc, phasespace in zip(self.subprocesses, self.phasespaces):
-            phasespace = subproc.build_madnis(phasespace)
-            if len(self.subprocesses) > 1:
-                status_func = lambda *args: self.update_madnis_status_multi(
-                    subproc.subproc_id, *args
-                )
-            else:
-                status_func = self.update_madnis_status_single
-            subproc.train_madnis(phasespace, status_func)
-            madnis_phasespaces.append(phasespace)
-        self.phasespaces = madnis_phasespaces
-        for context in self.contexts[1:]:
-            context.copy_globals_from(self.contexts[0])
-        self.event_generator = self.build_event_generator(madnis_phasespaces)
-
     def update_madnis_status_single(
         self, batch: int, batch_target: int, loss: float, lr: float, channel_count: int
     ) -> None:
diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml
index 28410c2626..b30b3b3f4a 100644
--- a/madgraph/iolibs/template_files/mg7/run_card.toml
+++ b/madgraph/iolibs/template_files/mg7/run_card.toml
@@ -121,6 +121,7 @@ lr_scheduler = %(madnis.lr_scheduler)s # options: none, cosine
 adam_beta1 = %(madnis.adam_beta1)s
 adam_beta2 = %(madnis.adam_beta2)s
 adam_eps = %(madnis.adam_eps)s
+grad_clip_threshold = %(madnis.grad_clip_threshold)s
 train_mcw = %(madnis.train_mcw)s
 buffer_capacity = %(madnis.buffer_capacity)s
 minimum_buffer_size = %(madnis.minimum_buffer_size)s
diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py
index 5d5b0c63dd..e608ea1637 100755
--- a/madgraph/various/banner.py
+++ b/madgraph/various/banner.py
@@ -6572,10 +6572,11 @@ def default_setup(self):
         self.add_toml_param('madnis', 'adam_beta1', 0.9)
         self.add_toml_param('madnis', 'adam_beta2', 0.999)
         self.add_toml_param('madnis', 'adam_eps', 1e-8)
+        self.add_toml_param('madnis', 'grad_clip_threshold', 0.0)
         self.add_toml_param('madnis', 'train_mcw', True)
-        self.add_toml_param('madnis', 'buffer_capacity', 0)
-        self.add_toml_param('madnis', 'minimum_buffer_size', 50000)
-        self.add_toml_param('madnis', 'buffered_steps', 0)
+        self.add_toml_param('madnis', 'buffer_capacity', 100000)
+        self.add_toml_param('madnis', 'minimum_buffer_size', 10000)
+        self.add_toml_param('madnis', 'buffered_steps', 5)
         self.add_toml_param('madnis', 'buffer_unweighting_quantile', 0.99)
         self.add_toml_param('madnis', 'uniform_channel_ratio', 0.1)
         self.add_toml_param('madnis', 'integration_history_length', 100)
diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt
index c7a6dd13db..6a58ae617a 100644
--- a/madspace/CMakeLists.txt
+++ b/madspace/CMakeLists.txt
@@ -173,6 +173,7 @@ add_library(
     src/driver/event_generator.cpp
     src/driver/channel_generator.cpp
     src/driver/generator_data.cpp
+    src/driver/status_file.cpp
     src/driver/thread_pool.cpp
     src/driver/io.cpp
     src/driver/vegas_optimizer.cpp
diff --git a/madspace/include/madspace/compgraphs/function_builder_mixin.inc b/madspace/include/madspace/compgraphs/function_builder_mixin.inc
index 360d8db14c..7b1aa8e3ef 100644
--- a/madspace/include/madspace/compgraphs/function_builder_mixin.inc
+++ b/madspace/include/madspace/compgraphs/function_builder_mixin.inc
@@ -87,6 +87,10 @@ Value reduce_sum_vector(Value in1) {
     return instruction("reduce_sum_vector", {in1})[0];
 }
 
+Value batch_reduce_sum(Value in1) {
+    return instruction("batch_reduce_sum", {in1})[0];
+}
+
 Value batch_reduce_mean(Value in1) {
     return instruction("batch_reduce_mean", {in1})[0];
 }
diff --git a/madspace/include/madspace/compgraphs/opcode_mixin.inc b/madspace/include/madspace/compgraphs/opcode_mixin.inc
index ba664afd04..88779728c2 100644
--- a/madspace/include/madspace/compgraphs/opcode_mixin.inc
+++ b/madspace/include/madspace/compgraphs/opcode_mixin.inc
@@ -19,139 +19,140 @@ mul = 17,
 div = 18,
 reduce_sum = 19,
 reduce_sum_vector = 20,
-batch_reduce_mean = 21,
-batch_reduce_mean_keepdim = 22,
-reduce_product = 23,
-sqrt = 24,
-square = 25,
-min = 26,
-max = 27,
-obs_sqrt_s = 28,
-obs_e = 29,
-obs_px = 30,
-obs_py = 31,
-obs_pz = 32,
-obs_mass = 33,
-obs_pt = 34,
-obs_p_mag = 35,
-obs_phi = 36,
-obs_theta = 37,
-obs_y = 38,
-obs_y_abs = 39,
-obs_eta = 40,
-obs_eta_abs = 41,
-obs_delta_eta = 42,
-obs_delta_phi = 43,
-obs_delta_r = 44,
-boost_beam = 45,
-boost_beam_inverse = 46,
-com_p_in = 47,
-r_to_x1x2 = 48,
-x1x2_to_r = 49,
-diff_cross_section = 50,
-two_body_decay_com = 51,
-two_body_decay_com_inverse = 52,
-two_body_decay = 53,
-two_body_decay_inverse = 54,
-two_to_two_particle_scattering_com = 55,
-two_to_two_particle_scattering_com_inverse = 56,
-two_to_two_particle_scattering = 57,
-two_to_two_particle_scattering_inverse = 58,
-two_to_three_particle_scattering = 59,
-two_to_three_particle_scattering_inverse = 60,
-double_t_scattering = 61,
-double_t_scattering_inverse = 62,
-three_body_decay_com = 63,
-three_body_decay_com_inverse = 64,
-three_body_decay = 65,
-three_body_decay_inverse = 66,
-t_inv_min_max = 67,
-t_inv_value_and_min_max = 68,
-t_inv_min_max_cut = 69,
-t_inv_value_and_min_max_cut = 70,
-t1_inv_min_max_doublet = 71,
-t1_inv_value_and_min_max_doublet = 72,
-t2_inv_min_max_doublet = 73,
-t2_inv_value_and_min_max_doublet = 74,
-s23_min_max = 75,
-s23_value_and_min_max = 76,
-s23_min_max_cut = 77,
-s23_value_and_min_max_cut = 78,
-invariants_from_momenta = 79,
-sde2_channel_weights = 80,
-subchannel_weights = 81,
-apply_subchannel_weights = 82,
-pt_eta_phi_x = 83,
-mirror_momenta = 84,
-momenta_to_x1x2 = 85,
-uniform_invariant = 86,
-uniform_invariant_inverse = 87,
-breit_wigner_invariant = 88,
-breit_wigner_invariant_inverse = 89,
-stable_invariant = 90,
-stable_invariant_inverse = 91,
-stable_invariant_nu = 92,
-stable_invariant_nu_inverse = 93,
-fast_rambo_massless = 94,
-fast_rambo_massless_inverse = 95,
-fast_rambo_massless_com = 96,
-fast_rambo_massive = 97,
-fast_rambo_massive_inverse = 98,
-fast_rambo_massive_com = 99,
-cut_unphysical = 100,
-cut_one = 101,
-cut_all = 102,
-cut_any = 103,
-scale_transverse_energy = 104,
-scale_transverse_mass = 105,
-scale_half_transverse_mass = 106,
-scale_partonic_energy = 107,
-chili_forward = 108,
-chili_inverse = 109,
-matrix_element = 110,
-collect_channel_weights = 111,
-interpolate_pdf = 112,
-interpolate_alpha_s = 113,
-matmul = 114,
-relu = 115,
-leaky_relu = 116,
-elu = 117,
-gelu = 118,
-sigmoid = 119,
-softplus = 120,
-rqs_reshape = 121,
-rqs_find_bin = 122,
-rqs_forward = 123,
-rqs_inverse = 124,
-softmax = 125,
-softmax_prior = 126,
-sample_discrete = 127,
-sample_discrete_inverse = 128,
-sample_discrete_probs = 129,
-sample_discrete_probs_inverse = 130,
-discrete_histogram = 131,
-permute_momenta = 132,
-gather = 133,
-gather_int = 134,
-gather_vector = 135,
-select_int = 136,
-select = 137,
-select_vector = 138,
-argsort = 139,
-quantile = 140,
-one_hot = 141,
-madnis_abs_weight = 142,
-madnis_softclip = 143,
-madnis_variance = 144,
-madnis_single_channel_variance = 145,
-madnis_multi_channel_variance = 146,
-nonzero = 147,
-batch_gather = 148,
-batch_scatter = 149,
-random = 150,
-random_int = 151,
-unweight = 152,
-vegas_forward = 153,
-vegas_inverse = 154,
-vegas_histogram = 155,
-histogram = 156
+batch_reduce_sum = 21,
+batch_reduce_mean = 22,
+batch_reduce_mean_keepdim = 23,
+reduce_product = 24,
+sqrt = 25,
+square = 26,
+min = 27,
+max = 28,
+obs_sqrt_s = 29,
+obs_e = 30,
+obs_px = 31,
+obs_py = 32,
+obs_pz = 33,
+obs_mass = 34,
+obs_pt = 35,
+obs_p_mag = 36,
+obs_phi = 37,
+obs_theta = 38,
+obs_y = 39,
+obs_y_abs = 40,
+obs_eta = 41,
+obs_eta_abs = 42,
+obs_delta_eta = 43,
+obs_delta_phi = 44,
+obs_delta_r = 45,
+boost_beam = 46,
+boost_beam_inverse = 47,
+com_p_in = 48,
+r_to_x1x2 = 49,
+x1x2_to_r = 50,
+diff_cross_section = 51,
+two_body_decay_com = 52,
+two_body_decay_com_inverse = 53,
+two_body_decay = 54,
+two_body_decay_inverse = 55,
+two_to_two_particle_scattering_com = 56,
+two_to_two_particle_scattering_com_inverse = 57,
+two_to_two_particle_scattering = 58,
+two_to_two_particle_scattering_inverse = 59,
+two_to_three_particle_scattering = 60,
+two_to_three_particle_scattering_inverse = 61,
+double_t_scattering = 62,
+double_t_scattering_inverse = 63,
+three_body_decay_com = 64,
+three_body_decay_com_inverse = 65,
+three_body_decay = 66,
+three_body_decay_inverse = 67,
+t_inv_min_max = 68,
+t_inv_value_and_min_max = 69,
+t_inv_min_max_cut = 70,
+t_inv_value_and_min_max_cut = 71,
+t1_inv_min_max_doublet = 72,
+t1_inv_value_and_min_max_doublet = 73,
+t2_inv_min_max_doublet = 74,
+t2_inv_value_and_min_max_doublet = 75,
+s23_min_max = 76,
+s23_value_and_min_max = 77,
+s23_min_max_cut = 78,
+s23_value_and_min_max_cut = 79,
+invariants_from_momenta = 80,
+sde2_channel_weights = 81,
+subchannel_weights = 82,
+apply_subchannel_weights = 83,
+pt_eta_phi_x = 84,
+mirror_momenta = 85,
+momenta_to_x1x2 = 86,
+uniform_invariant = 87,
+uniform_invariant_inverse = 88,
+breit_wigner_invariant = 89,
+breit_wigner_invariant_inverse = 90,
+stable_invariant = 91,
+stable_invariant_inverse = 92,
+stable_invariant_nu = 93,
+stable_invariant_nu_inverse = 94,
+fast_rambo_massless = 95,
+fast_rambo_massless_inverse = 96,
+fast_rambo_massless_com = 97,
+fast_rambo_massive = 98,
+fast_rambo_massive_inverse = 99,
+fast_rambo_massive_com = 100,
+cut_unphysical = 101,
+cut_one = 102,
+cut_all = 103,
+cut_any = 104,
+scale_transverse_energy = 105,
+scale_transverse_mass = 106,
+scale_half_transverse_mass = 107,
+scale_partonic_energy = 108,
+chili_forward = 109,
+chili_inverse = 110,
+matrix_element = 111,
+collect_channel_weights = 112,
+interpolate_pdf = 113,
+interpolate_alpha_s = 114,
+matmul = 115,
+relu = 116,
+leaky_relu = 117,
+elu = 118,
+gelu = 119,
+sigmoid = 120,
+softplus = 121,
+rqs_reshape = 122,
+rqs_find_bin = 123,
+rqs_forward = 124,
+rqs_inverse = 125,
+softmax = 126,
+softmax_prior = 127,
+sample_discrete = 128,
+sample_discrete_inverse = 129,
+sample_discrete_probs = 130,
+sample_discrete_probs_inverse = 131,
+discrete_histogram = 132,
+permute_momenta = 133,
+gather = 134,
+gather_int = 135,
+gather_vector = 136,
+select_int = 137,
+select = 138,
+select_vector = 139,
+argsort = 140,
+quantile = 141,
+one_hot = 142,
+madnis_abs_weight = 143,
+madnis_softclip = 144,
+madnis_variance = 145,
+madnis_single_channel_variance = 146,
+madnis_multi_channel_variance = 147,
+nonzero = 148,
+batch_gather = 149,
+batch_scatter = 150,
+random = 151,
+random_int = 152,
+unweight = 153,
+vegas_forward = 154,
+vegas_inverse = 155,
+vegas_histogram = 156,
+histogram = 157
diff --git a/madspace/include/madspace/driver.hpp b/madspace/include/madspace/driver.hpp
index 1982b36141..d17ce901d0 100644
--- a/madspace/include/madspace/driver.hpp
+++ b/madspace/include/madspace/driver.hpp
@@ -10,6 +10,7 @@
 #include "driver/lhe_output.hpp"
 #include "driver/logger.hpp"
 #include "driver/madnis_training.hpp"
+#include "driver/status_file.hpp"
 #include "driver/tensor.hpp"
 #include "driver/thread_pool.hpp"
 #include "driver/vegas_optimizer.hpp"
diff --git a/madspace/include/madspace/driver/adam_optimizer.hpp b/madspace/include/madspace/driver/adam_optimizer.hpp
index c8df5138c6..fafb83749f 100644
--- a/madspace/include/madspace/driver/adam_optimizer.hpp
+++ b/madspace/include/madspace/driver/adam_optimizer.hpp
@@ -2,9 +2,22 @@
 
 #include "madspace/driver/backend.hpp"
 #include "madspace/driver/context.hpp"
+#include "madspace/phasespace/base.hpp"
 
 namespace madspace {
 
+class GradientClipper : public FunctionGenerator {
+public:
+    GradientClipper(double threshold);
+
+private:
+    NamedVector<Value> build_function_impl(
+        FunctionBuilder& fb, const NamedVector<Value>& args
+    ) const override;
+
+    double _threshold;
+};
+
 class AdamOptimizer {
 public:
     enum LRSchedule {
@@ -20,7 +33,8 @@ class AdamOptimizer {
         std::size_t step_count = 0,
         double beta1 = 0.9,
         double beta2 = 0.999,
-        double eps = 1e-8
+        double eps = 1e-8,
+        double grad_clip_threshold = 0.0
     );
     TensorVec step(const TensorVec& inputs);
     void replace_function(const Function& function);
@@ -40,6 +54,7 @@ class AdamOptimizer {
     double _beta1;
     double _beta2;
     double _eps;
+    RuntimePtr _grad_clipper;
     Tensor _one;
     Tensor _parameter;
     Tensor _exp_avg;
diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp
index 0e9428dff6..ebc4cd3a24 100644
--- a/madspace/include/madspace/driver/event_generator.hpp
+++ b/madspace/include/madspace/driver/event_generator.hpp
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <chrono>
+#include <memory>
 #include <optional>
 #include <random>
 #include <vector>
@@ -15,6 +16,7 @@
 #include "madspace/driver/generator_data.hpp"
 #include "madspace/driver/io.hpp"
 #include "madspace/driver/lhe_output.hpp"
+#include "madspace/driver/status_file.hpp"
 #include "madspace/driver/vegas_optimizer.hpp"
 #include "madspace/phasespace.hpp"
 
@@ -30,7 +32,7 @@ class EventGenerator {
     EventGenerator(
         const std::vector<ContextPtr>& contexts,
         const std::vector<std::shared_ptr<ChannelEventGenerator>>& channels,
-        const std::string& status_file = "",
+        std::shared_ptr<StatusFile> status_file = nullptr,
         const GeneratorConfig& config = default_config
     );
     void survey();
@@ -38,7 +40,8 @@ class EventGenerator {
     void combine_to_compact_npy(const std::string& file_name);
     void combine_to_lhe_npy(const std::string& file_name, LHECompleter& lhe_completer);
     void combine_to_lhe(
-        const std::string& file_name, LHECompleter& lhe_completer,
+        const std::string& file_name,
+        LHECompleter& lhe_completer,
         const LHEMeta& meta = {}
     );
     GeneratorStatus status() const { return _status; }
@@ -78,10 +81,9 @@ class EventGenerator {
     std::chrono::time_point<std::chrono::steady_clock> _start_time;
     std::size_t _start_cpu_microsec;
     std::chrono::time_point<std::chrono::steady_clock> _last_print_time;
-    std::chrono::time_point<std::chrono::steady_clock> _last_status_time;
     PrettyBox _pretty_box_upper;
     PrettyBox _pretty_box_lower;
-    std::string _status_file;
+    std::shared_ptr<StatusFile> _status_file;
     std::unordered_map<std::string, TimingData> _timing_data;
 
     bool start_jobs();
diff --git a/madspace/include/madspace/driver/madnis_training.hpp b/madspace/include/madspace/driver/madnis_training.hpp
index 9327601197..ce5010a7c3 100644
--- a/madspace/include/madspace/driver/madnis_training.hpp
+++ b/madspace/include/madspace/driver/madnis_training.hpp
@@ -1,11 +1,13 @@
 #pragma once
 
 #include <chrono>
+#include <memory>
 
 #include "madspace/compgraphs.hpp"
 #include "madspace/driver/adam_optimizer.hpp"
 #include "madspace/driver/format.hpp"
 #include "madspace/driver/logger.hpp"
+#include "madspace/driver/status_file.hpp"
 #include "madspace/phasespace.hpp"
 
 namespace madspace {
@@ -35,6 +37,7 @@ class MadnisTraining {
         double adam_beta1 = 0.9;
         double adam_beta2 = 0.999;
         double adam_eps = 1e-8;
+        double grad_clip_threshold = 0.0;
         std::size_t buffer_capacity = 0;
         std::size_t minimum_buffer_size = 10000;
         std::size_t buffered_steps = 0;
@@ -53,6 +56,30 @@ class MadnisTraining {
     std::vector<std::size_t> active_channels() const;
     std::size_t active_channel_count() const { return _channels.size(); }
     double average_loss() const;
+    double average_learning_rate() const;
+    double buffered_fraction() const;
+    std::size_t generated_event_count() const { return _generated_event_count; }
+    std::size_t buffer_event_count() const;
+
+    // Status histories, one entry appended every config.log_interval batches,
+    // for reporting training progress (e.g. to a StatusFile).
+    const std::vector<std::size_t>& status_batches() const { return _status_batches; }
+    const std::vector<double>& status_losses() const { return _status_losses; }
+    const std::vector<std::size_t>& status_channel_counts() const {
+        return _status_channel_counts;
+    }
+    const std::vector<double>& status_learning_rates() const {
+        return _status_learning_rates;
+    }
+    const std::vector<double>& status_buffered_fractions() const {
+        return _status_buffered_fractions;
+    }
+    const std::vector<std::size_t>& status_generated_events() const {
+        return _status_generated_events;
+    }
+    const std::vector<std::size_t>& status_buffer_sizes() const {
+        return _status_buffer_sizes;
+    }
 
 private:
     struct SampleBatch {
@@ -93,8 +120,12 @@ class MadnisTraining {
     TensorVec build_buffered_training_batch(const std::vector<size_t>& counts);
     void process_job_results(const std::vector<std::size_t>& job_ids);
     void buffer_store(ChannelData& channel, SampleBatch& samples);
-    void
-    update_history(const TensorVec& results, const std::vector<std::size_t>& counts);
+    void update_history(
+        const TensorVec& results,
+        const std::vector<std::size_t>& counts,
+        double learning_rate,
+        bool buffered
+    );
     void drop_channels();
     void freeze_cwnet();
 
@@ -109,7 +140,17 @@ class MadnisTraining {
     std::vector<ChannelData> _channels;
     std::unordered_map<std::size_t, SampleJob> _running_jobs;
     std::vector<double> _loss_history;
+    std::vector<double> _lr_history;
+    std::vector<bool> _buffered_history;
     std::size_t _loss_history_index = 0;
+    std::vector<std::size_t> _status_batches;
+    std::vector<double> _status_losses;
+    std::vector<std::size_t> _status_channel_counts;
+    std::vector<double> _status_learning_rates;
+    std::vector<double> _status_buffered_fractions;
+    std::vector<std::size_t> _status_generated_events;
+    std::vector<std::size_t> _status_buffer_sizes;
+    std::size_t _generated_event_count = 0;
     std::size_t _job_id = 0;
     Tensor _generator_params;
     std::vector<std::size_t> _arg_permutation;
@@ -124,7 +165,8 @@ class MultiMadnisTraining {
         ContextPtr optimizer_context,
         const MadnisTraining::Config& config,
         const nested_vector2<std::shared_ptr<Integrand>>& integrands,
-        const std::vector<std::optional<ChannelWeightNetwork>>& cwnets
+        const std::vector<std::optional<ChannelWeightNetwork>>& cwnets,
+        std::shared_ptr<StatusFile> status_file = nullptr
     );
     void train();
     nested_vector2<std::size_t> active_channels() const;
@@ -137,6 +179,7 @@ class MultiMadnisTraining {
         double loss,
         std::size_t chan_count
     );
+    void write_status(std::size_t subproc_index, std::size_t batch_index, bool done);
 
     MadnisTraining::Config _config;
     std::vector<MadnisTraining> _subprocesses;
@@ -145,6 +188,7 @@ class MultiMadnisTraining {
     std::chrono::time_point<std::chrono::steady_clock> _last_print_time;
     PrettyBox _pretty_box_upper;
     PrettyBox _pretty_box_lower;
+    std::shared_ptr<StatusFile> _status_file;
 };
 
 } // namespace madspace
diff --git a/madspace/include/madspace/driver/status_file.hpp b/madspace/include/madspace/driver/status_file.hpp
new file mode 100644
index 0000000000..3cea8d1751
--- /dev/null
+++ b/madspace/include/madspace/driver/status_file.hpp
@@ -0,0 +1,36 @@
+#pragma once
+
+#include <chrono>
+#include <string>
+
+#include <nlohmann/json.hpp>
+
+namespace madspace {
+
+// Accumulates JSON status content from possibly multiple sources and
+// periodically rewrites it to a file.
+class StatusFile {
+public:
+    // file_name is where the status is written; writes happen at most every
+    // min_interval_sec, unless forced (see write()).
+    explicit StatusFile(const std::string& file_name, double min_interval_sec = 10.0);
+
+    // Merges content into the accumulated status (top-level keys are
+    // overwritten, "run_times" is merged one level deeper) and rewrites the
+    // file if force_write is set, if this is the first write, or if
+    // min_interval_sec has passed since the last write.
+    void write(const nlohmann::json& content, bool force_write = false);
+
+    // If a status was ever written and its "status" field is not "done",
+    // marks it "done" and writes it out one last time.
+    ~StatusFile();
+
+private:
+    std::string _file_name;
+    std::chrono::duration<double> _min_interval;
+    std::chrono::time_point<std::chrono::steady_clock> _last_write_time;
+    bool _has_written = false;
+    nlohmann::json _content = nlohmann::json::object();
+};
+
+} // namespace madspace
diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml
index cc80553819..fbade8bca7 100644
--- a/madspace/instruction_set.yaml
+++ b/madspace/instruction_set.yaml
@@ -253,6 +253,18 @@ reduce_sum_vector:
       desc:
   dims: 0
 
+batch_reduce_sum:
+  inputs:
+    - name: in1
+      type: [float]
+      desc:
+  outputs:
+    - name: out
+      type: [float, single]
+      desc:
+  differentiable: True
+  custom_op: True
+
 batch_reduce_mean:
   inputs:
     - name: in1
diff --git a/madspace/madspace/_madspace_py.pyi b/madspace/madspace/_madspace_py.pyi
index d64a398e7f..0faed1b229 100644
--- a/madspace/madspace/_madspace_py.pyi
+++ b/madspace/madspace/_madspace_py.pyi
@@ -165,6 +165,7 @@ class AdamOptimizer:
         beta1: typing.SupportsFloat = 0.9,
         beta2: typing.SupportsFloat = 0.999,
         eps: typing.SupportsFloat = 1e-08,
+        grad_clip_threshold: typing.SupportsFloat = 0.0,
     ) -> None: ...
     def context(self) -> Context: ...
     def input_types(self) -> list[Type]: ...
@@ -581,7 +582,9 @@ class EventGenerator:
     def channel_status(self) -> list[GeneratorStatus]: ...
     def channels(self) -> list[ChannelEventGenerator]: ...
     def combine_to_compact_npy(self, file_name: str) -> None: ...
-    def combine_to_lhe(self, file_name: str, lhe_completer: LHECompleter, meta: LHEMeta = ...) -> None: ...
+    def combine_to_lhe(
+        self, file_name: str, lhe_completer: LHECompleter, meta: LHEMeta = ...
+    ) -> None: ...
     def combine_to_lhe_npy(
         self, file_name: str, lhe_completer: LHECompleter
     ) -> None: ...
@@ -1735,6 +1738,10 @@ class MadnisConfig:
     @gpu_generator_batch_size.setter
     def gpu_generator_batch_size(self, arg0: typing.SupportsInt) -> None: ...
     @property
+    def grad_clip_threshold(self) -> float: ...
+    @grad_clip_threshold.setter
+    def grad_clip_threshold(self, arg0: typing.SupportsFloat) -> None: ...
+    @property
     def integration_history_length(self) -> int: ...
     @integration_history_length.setter
     def integration_history_length(self, arg0: typing.SupportsInt) -> None: ...
diff --git a/madspace/madspace/madnis/__init__.py b/madspace/madspace/madnis/__init__.py
deleted file mode 100644
index 31f5dbab3c..0000000000
--- a/madspace/madspace/madnis/__init__.py
+++ /dev/null
@@ -1,44 +0,0 @@
-"""
-This module contains functions and classes to train neural importance sampling networks and
-evaluate the integration and sampling performance.
-"""
-
-from .buffer import Buffer
-from .channel_grouping import ChannelData, ChannelGroup, ChannelGrouping
-from .distribution import Distribution
-from .integrand import Integrand
-from .integrator import Integrator, SampleBatch, TrainingStatus
-from .interface import (
-    IntegrandDistribution,
-    IntegrandFunction,
-    build_madnis_integrand,
-)
-from .losses import (
-    kl_divergence,
-    kl_divergence_softclip,
-    multi_channel_loss,
-    stratified_variance,
-    stratified_variance_softclip,
-    variance,
-)
-
-__all__ = [
-    "Integrator",
-    "TrainingStatus",
-    "SampleBatch",
-    "Integrand",
-    "Buffer",
-    "multi_channel_loss",
-    "stratified_variance",
-    "stratified_variance_softclip",
-    "variance",
-    "kl_divergence",
-    "kl_divergence_softclip",
-    "ChannelGroup",
-    "ChannelData",
-    "ChannelGrouping",
-    "Distribution",
-    "IntegrandDistribution",
-    "IntegrandFunction",
-    "build_madnis_integrand",
-]
diff --git a/madspace/madspace/madnis/buffer.py b/madspace/madspace/madnis/buffer.py
deleted file mode 100644
index d481c7d6de..0000000000
--- a/madspace/madspace/madnis/buffer.py
+++ /dev/null
@@ -1,167 +0,0 @@
-from collections.abc import Callable
-
-import torch
-import torch.nn as nn
-
-
-class Buffer(nn.Module):
-    """
-    Circular buffer for multiple tensors with different shapes. The class is a torch.nn.Module to
-    allow for simple storage.
-    """
-
-    def __init__(
-        self,
-        capacity: int,
-        shapes: list[tuple[int, ...]],
-        persistent: bool = True,
-        dtypes: list[torch.dtype | None] | None = None,
-    ):
-        """
-        Args:
-            capacity: maximum number of samples stored in the buffer
-            shapes: shapes of the tensors to be stored, without batch dimension. If a shape is
-                None, no tensor is stored at that position. This allows for simpler handling of
-                optional stored fields.
-            persistent: if True, the content of the buffer is part of the module's state_dict
-            dtypes: if different from None, specifies the tensors which have a non-standard dtype
-        """
-        super().__init__()
-        self.keys = []
-        if dtypes is None:
-            dtypes = [None] * len(shapes)
-        for i, (shape, dtype) in enumerate(zip(shapes, dtypes)):
-            key = f"buffer{i}"
-            self.register_buffer(
-                key,
-                None if shape is None else torch.zeros((capacity, *shape), dtype=dtype),
-                persistent,
-            )
-            self.keys.append(key)
-        self.capacity = capacity
-        self.size = 0
-        self.store_index = 0
-
-    def _batch_slices(self, batch_size: int) -> slice:
-        """
-        Returns slices that split up the buffer into batches of at most ``batch_size``, respecting
-        the buffer size and periodic boundary.
-        """
-        start = self.store_index
-        while start < self.size:
-            stop = min(start + batch_size, self.size)
-            yield slice(start, stop)
-            start = stop
-        start = 0
-        while start < self.store_index:
-            stop = min(start + batch_size, self.store_index)
-            yield slice(start, stop)
-            start = stop
-
-    def _buffer_fields(self) -> torch.Tensor | None:
-        """
-        Iterates over the buffered tensors, without removing the padding if the buffer is not full.
-
-        Returns:
-            The buffered tensors, or None if a tensor was initialized with shape None
-        """
-        for key in self.keys:
-            yield getattr(self, key)
-
-    def __iter__(self) -> torch.Tensor | None:
-        """
-        Iterates over the buffered tensors
-
-        Returns:
-            The buffered tensors, or None if a tensor was initialized with shape None
-        """
-        for key in self.keys:
-            buffer = getattr(self, key)
-            yield None if buffer is None else buffer[: self.size]
-
-    def store(self, *tensors: torch.Tensor | None):
-        """
-        Adds the given tensors to the buffer. If the buffer is full, the oldest stored samples are
-        overwritten.
-
-        Args:
-            tensors: samples to be stored. The shapes of the tensors after the batch dimension must
-                match the shapes given during initialization. The argument can be None if the
-                corresponding shape was None during initialization.
-        """
-        store_slice1 = None
-        for buffer, data in zip(self._buffer_fields(), tensors):
-            if data is None:
-                continue
-            if store_slice1 is None:
-                size = min(data.shape[0], self.capacity)
-                end_index = self.store_index + size
-                if end_index < self.capacity:
-                    store_slice1 = slice(self.store_index, end_index)
-                    store_slice2 = slice(0, 0)
-                    load_slice1 = slice(0, size)
-                    load_slice2 = slice(0, 0)
-                else:
-                    store_slice1 = slice(self.store_index, self.capacity)
-                    store_slice2 = slice(0, end_index - self.capacity)
-                    load_slice1 = slice(0, self.capacity - self.store_index)
-                    load_slice2 = slice(self.capacity - self.store_index, size)
-                self.store_index = end_index % self.capacity
-                self.size = min(self.size + size, self.capacity)
-            buffer[store_slice1] = data[load_slice1]
-            buffer[store_slice2] = data[load_slice2]
-
-    def filter(
-        self,
-        predicate: Callable[[tuple[torch.Tensor | None, ...]], torch.Tensor],
-        batch_size: int = 100000,
-    ):
-        """
-        Removes samples from the buffer that do not fulfill the criterion given by the predicate
-        function.
-
-        Args:
-            predicate: function that returns a mask for a batch of samples, given a tuple with
-                all the buffered fields as argument
-            batch_size: maximal batch size to limit memory usage
-        """
-        masks = []
-        masked_size = 0
-        for batch_slice in self._batch_slices(batch_size):
-            mask = predicate(
-                tuple(
-                    None if t is None else t[batch_slice] for t in self._buffer_fields()
-                )
-            )
-            masked_size += torch.count_nonzero(mask)
-            masks.append(mask)
-        for buffer in self._buffer_fields():
-            if buffer is None:
-                continue
-            buffer[:masked_size] = torch.cat(
-                [
-                    buffer[batch_slice][mask]
-                    for batch_slice, mask in zip(self._batch_slices(batch_size), masks)
-                ],
-                dim=0,
-            )
-        self.size = masked_size
-        self.store_index = masked_size % self.capacity
-
-    def sample(self, count: int) -> list[torch.Tensor | None]:
-        """
-        Returns a batch of samples drawn from the buffer without replacement.
-
-        Args:
-            count: number of samples
-        Returns:
-            samples drawn from the buffer
-        """
-        weights = next(b for b in self._buffer_fields() if b is not None).new_ones(
-            self.size
-        )
-        indices = torch.multinomial(weights, min(count, self.size), replacement=False)
-        return [
-            None if buffer is None else buffer[indices]
-            for buffer in self._buffer_fields()
-        ]
diff --git a/madspace/madspace/madnis/channel_grouping.py b/madspace/madspace/madnis/channel_grouping.py
deleted file mode 100644
index 76ee970135..0000000000
--- a/madspace/madspace/madnis/channel_grouping.py
+++ /dev/null
@@ -1,85 +0,0 @@
-from dataclasses import dataclass
-
-
-@dataclass
-class ChannelGroup:
-    """
-    A group of channels
-
-    Args:
-        group_index: index of the group in the list of groups
-        target_index: index of the channel that all other channels in the group are mapped to
-        channel_indices: indices of the channels in the group
-    """
-
-    group_index: int
-    target_index: int
-    channel_indices: list[int]
-
-
-@dataclass
-class ChannelData:
-    """
-    Information about a single channel
-
-    Args:
-        channel_index: index of the channel
-        target_index: index of the channel that it is mapped to
-        group: channel group that the channel belongs to
-        remapped: True if the channel is remapped to another channel
-        position_in_group: index of the channel within its group
-    """
-
-    channel_index: int
-    target_index: int
-    group: ChannelGroup
-    remapped: bool
-    position_in_group: int
-
-
-class ChannelGrouping:
-    """
-    Class that encodes how channels are grouped together for a multi-channel integrand
-    """
-
-    def __init__(self, channel_assignment: list[int | None]):
-        """
-        Args:
-            channel_assignment: list with an entry for each channel. If None, the channel is not
-                remapped. Otherwise, the index of the channel to which it is mapped.
-        """
-        group_dict = {}
-        for source_channel, target_channel in enumerate(channel_assignment):
-            if target_channel is None:
-                group_dict[source_channel] = ChannelGroup(
-                    group_index=len(group_dict),
-                    target_index=source_channel,
-                    channel_indices=[source_channel],
-                )
-
-        self.channels: list[ChannelData] = []
-        self.groups: list[ChannelGroup] = list(group_dict.values())
-
-        for source_channel, target_channel in enumerate(channel_assignment):
-            if target_channel is None:
-                self.channels.append(
-                    ChannelData(
-                        channel_index=source_channel,
-                        target_index=source_channel,
-                        group=group_dict[source_channel],
-                        remapped=False,
-                        position_in_group=0,
-                    )
-                )
-            else:
-                group = group_dict[target_channel]
-                self.channels.append(
-                    ChannelData(
-                        channel_index=source_channel,
-                        target_index=target_channel,
-                        group=group,
-                        remapped=True,
-                        position_in_group=len(group.channel_indices),
-                    )
-                )
-                group.channel_indices.append(source_channel)
diff --git a/madspace/madspace/madnis/distribution.py b/madspace/madspace/madnis/distribution.py
deleted file mode 100644
index 120aad48dc..0000000000
--- a/madspace/madspace/madnis/distribution.py
+++ /dev/null
@@ -1,103 +0,0 @@
-import math
-from collections.abc import Callable
-from typing import Literal, Protocol
-
-import numpy as np
-import torch
-import torch.nn as nn
-
-L2PI = -0.5 * math.log(2 * math.pi)
-
-Mapping = Callable[[torch.Tensor, bool], tuple[torch.Tensor, torch.Tensor]]
-
-
-class Distribution(Protocol):
-    """
-    Protocol for a (potentially learnable) distribution that can be used for sampling and
-    density estimation, like a normalizing flow.
-    """
-
-    def sample(
-        self,
-        n: int,
-        c: torch.Tensor | None = None,
-        channel: torch.Tensor | list[int] | int | None = None,
-        return_log_prob: bool = False,
-        return_prob: bool = False,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
-        """
-        Draws samples following the distribution
-
-        Args:
-            n: number of samples
-            c: condition, shape (n, dims_c) or None for an unconditional flow
-            channel: encodes the channel of the samples. It must have one of the following types:
-
-                - ``Tensor``: integer tensor of shape (n, ), containing the channel index for every
-                  input sample;
-                - ``list``: list of integers, specifying the number of samples in each channel;
-                - ``int``: integer specifying a single channel containing all the samples;
-                - ``None``: used in the single-channel case or to indicate that all channels contain
-                  the same number of samples in the multi-channel case.
-            return_log_prob: if True, also return the log-probabilities
-            return_prob: if True, also return the probabilities
-            device: device of the returned tensor. Only required if no condition is given.
-            dtype: dtype of the returned tensor. Only required if no condition is given.
-        Returns:
-            samples with shape (n, dims_in). Depending on the arguments ``return_log_prob``,
-            ``return_prob``, this function should also return the log-probabilities with shape (n, ),
-            the probabilities with shape (n, ).
-        """
-        ...
-
-    def log_prob(
-        self,
-        x: torch.Tensor,
-        c: torch.Tensor | None = None,
-        channel: torch.Tensor | list[int] | int | None = None,
-    ) -> torch.Tensor:
-        """
-        Computes the log-probabilities of the input data.
-
-        Args:
-            x: input data, shape (n, dims_in)
-            c: condition, shape (n, dims_c) or None for an unconditional flow
-            channel: encodes the channel of the samples. It must have one of the following types:
-
-                - ``Tensor``: integer tensor of shape (n, ), containing the channel index for every
-                  input sample;
-                - ``list``: list of integers, specifying the number of samples in each channel;
-                - ``int``: integer specifying a single channel containing all the samples;
-                - ``None``: used in the single-channel case or to indicate that all channels contain
-                  the same number of samples in the multi-channel case.
-        Returns:
-            log-probabilities with shape (n, )
-        """
-        return self.prob(x, c, channel).log()
-
-    def prob(
-        self,
-        x: torch.Tensor,
-        c: torch.Tensor | None = None,
-        channel: torch.Tensor | list[int] | int | None = None,
-    ) -> torch.Tensor:
-        """
-        Computes the probabilities of the input data.
-
-        Args:
-            x: input data, shape (n, dims_in)
-            c: condition, shape (n, dims_c) or None for an unconditional flow
-            channel: encodes the channel of the samples. It must have one of the following types:
-
-                - ``Tensor``: integer tensor of shape (n, ), containing the channel index for every
-                  input sample;
-                - ``list``: list of integers, specifying the number of samples in each channel;
-                - ``int``: integer specifying a single channel containing all the samples;
-                - ``None``: used in the single-channel case or to indicate that all channels contain
-                  the same number of samples in the multi-channel case.
-        Returns:
-            probabilities with shape (n, )
-        """
-        return self.log_prob(x, c, channel).exp()
diff --git a/madspace/madspace/madnis/integrand.py b/madspace/madspace/madnis/integrand.py
deleted file mode 100644
index 41fe917e12..0000000000
--- a/madspace/madspace/madnis/integrand.py
+++ /dev/null
@@ -1,175 +0,0 @@
-from collections.abc import Callable
-from typing import Literal
-
-import torch
-import torch.nn as nn
-
-from .channel_grouping import ChannelGrouping
-
-
-class Integrand(nn.Module):
-    """
-    Class that wraps an integrand function and meta-data necessary to use advanced MadNIS
-    features like learnable multi-channel weights, grouped channels and channel weight priors.
-    """
-
-    def __init__(
-        self,
-        function: Callable,
-        input_dim: int,
-        bounds: list[list[float]] | None = None,
-        channel_count: int | None = None,
-        remapped_dim: int | None = None,
-        has_channel_weight_prior: bool = False,
-        channel_grouping: ChannelGrouping | None = None,
-        function_includes_sampling: bool = False,
-        update_active_channels_mask: Callable[[torch.Tensor], None] | None = None,
-        discrete_dims: list[int] = [],
-        discrete_dims_position: Literal["first", "last"] = "first",
-        discrete_prior_prob_function: (
-            Callable[[torch.Tensor, int], torch.Tensor] | None
-        ) = None,
-    ):
-        """
-        Args:
-            function: integrand function.
-                The signature depends on the other arguments:
-
-                  - single-channel integration, ``channel_count=None``: ``x -> f``
-                  - basic multi-channel integration, ``remapped_dim=None``,
-                    ``has_channel_weight_prior=False``: ``(x, c) -> f``
-                  - with channel weights, ``remapped_dim=None``, ``has_channel_weight_prior=True``:
-                    ``(x, c) -> (f, alpha)`` (no trainable channel weights possible)
-                  - with channel-dependent mapping, ``remapped_dim: int``,
-                    ``has_channel_weight_prior=False``: ``(x, c) -> (f, y)``
-                  - all features, ``remapped_dim: int``, ``has_channel_weight_prior=True``:
-                    ``(x, c) -> (f, y, alpha)``
-
-                with the following tensors:
-
-                  - ``x`` is a point generated by the importance sampling, shape (n, input_dim),
-                  - ``c`` is the channel index, shape (n, ),
-                  - ``f`` is the integrand value, shape (n, ),
-                  - ``y`` is the point after applying a channel-dependent mapping, shape
-                    (n, remapped_dim)
-                  - ``alpha`` is the prior channel weight, shape (n, channel_count).
-            input_dim: dimension of the integration space
-            bounds: List of pairs ``[lower bound, upper bound]`` of the integration interval for
-                all dimensions. The integrand is rescaled so that the MadNIS training can be
-                performed on the unit hypercube. If None, the unit hypercube is used as integration
-                domain.
-            channel_count: None in the single-channel case, specifies the number of channels
-                otherwise.
-            remapped_dim: If different from None, it gives the dimension of a remapped space,
-                with a channel-dependent mapping computed as part of the integrand function.
-            has_channel_weight_prior: If True, the integrand returns channel weights
-            channel_grouping: ChannelGrouping object or None if all channels are independent
-        """
-        # TODO: update documentation
-        super().__init__()
-        self.input_dim = input_dim
-        self.remapped_dim = input_dim if remapped_dim is None else remapped_dim
-        self.channel_count = channel_count
-        self.has_channel_weight_prior = has_channel_weight_prior
-        self.channel_grouping = channel_grouping
-        self.function_includes_sampling = function_includes_sampling
-        self.update_active_channels_mask_func = update_active_channels_mask
-
-        self.discrete_dims = discrete_dims
-        self.discrete_dims_position = discrete_dims_position
-        self.discrete_prior_prob_function = discrete_prior_prob_function
-
-        if function_includes_sampling:
-            self.function = function
-        elif channel_count is None:
-            self.function = lambda x, channels: (function(x), None, None)
-        elif remapped_dim is None:
-            if has_channel_weight_prior:
-
-                def func(x, channels):
-                    w, prior = function(x, channels)
-                    return w, None, prior
-
-                self.function = func
-            else:
-                self.function = lambda x, channels: (function(x, channels), None, None)
-
-        elif has_channel_weight_prior:
-            self.function = function
-        else:
-
-            def func(x, channels):
-                w, y = function(x, channels)
-                return w, y, None
-
-            self.function = func
-
-        if bounds is not None:
-            bounds = torch.tensor(bounds)
-            self.register_buffer("scale", bounds[:, 1] - bounds[:, 0])
-            self.register_buffer("offset", bounds[:, 0])
-            self.register_buffer("scale_det", self.scale.prod())
-            old_func = self.function
-
-            def rescaled_func(x, channels):
-                w, y, prior = old_func(self.scale * x + self.offset, channels)
-                return self.scale_det * w, y, prior
-
-            self.function = rescaled_func
-
-        self.register_buffer(
-            "channel_id_map",
-            (
-                None
-                if self.channel_grouping is None
-                else torch.tensor(
-                    [
-                        channel.group.group_index
-                        for channel in self.channel_grouping.channels
-                    ]
-                )
-            ),
-        )
-
-    def unique_channel_count(self) -> int:
-        """
-        Returns the number of channels, or, if some channels are grouped together, the number of
-        channel groups
-        """
-        if self.channel_grouping is None:
-            return self.channel_count
-        else:
-            return len(self.channel_grouping.groups)
-
-    def remap_channels(self, channels: torch.Tensor | int) -> torch.Tensor | int:
-        """
-        Remaps channel indices to the indices of their respective channel groups if a
-        ``ChannelGrouping`` object was provided, otherwise returns the indices unchanged.
-
-        Args:
-            channels: channel indices, tensor with shape (n, ) or integer
-        Returns:
-            remapped channel indices, tensor with shape (n, ) or integer
-        """
-        if self.channel_grouping is None:
-            return channels
-        elif isinstance(channels, int):
-            return self.channel_id_map[channels].item()
-        else:
-            return self.channel_id_map[channels]
-
-    def update_active_channels_mask(self, mask: torch.Tensor) -> None:
-        if self.update_active_channels_mask_func is None:
-            return
-
-        full_mask = mask[
-            self.remap_channels(
-                torch.arange(len(self.channel_grouping.channels), device=mask.device)
-            )
-        ]
-        self.update_active_channels_mask_func(full_mask)
-
-    def forward(
-        self, x: torch.Tensor, channels: torch.Tensor | None
-    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
-        return self.function(x, channels)
diff --git a/madspace/madspace/madnis/integrator.py b/madspace/madspace/madnis/integrator.py
deleted file mode 100644
index c7b1729e40..0000000000
--- a/madspace/madspace/madnis/integrator.py
+++ /dev/null
@@ -1,989 +0,0 @@
-import itertools
-import signal
-import warnings
-from collections.abc import Callable, Iterable
-from dataclasses import astuple, dataclass
-from typing import Any, Literal
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.optim import Optimizer
-from torch.optim.lr_scheduler import LRScheduler
-
-from .buffer import Buffer
-from .distribution import Distribution
-from .integrand import Integrand
-from .losses import MultiChannelLoss, kl_divergence, stratified_variance, variance
-
-
-@dataclass
-class TrainingStatus:
-    """
-    Contains the MadNIS training status to pass it to a callback function.
-
-    Args:
-        step: optimization step
-        loss: loss from the optimization step
-        buffered: whether the optimization was performed on buffered samples
-        learning_rate: current learning rate if learning rate scheduler is present
-        dropped_channels: number of channels dropped after this optimization step
-    """
-
-    step: int
-    loss: float
-    buffered: bool
-    learning_rate: float | None
-    dropped_channels: int
-
-
-@dataclass
-class SampleBatch:
-    """
-    Contains a batch of samples
-
-    Args:
-        x: samples generated by the flow, shape (n, dim)
-        y: remapped samples returned by the integrand, shape (n, remapped_dim)
-        q_sample: probabilities of the samples, shape (n, )
-        func_vals: integrand value, shape (n, )
-        channels: channels indices for multi-channel integration, shape (n, ), otherwise None
-        alphas_prior: prior channel weights, shape (n, channels), or None for single-channel
-            integration
-        alpha_channel_indices: channel indices if not all prior channel weights are stored,
-            otherwise None
-        integration_channels: index of the channel group in case the integration is performed at the
-            level of channel groups, shape (n, ), otherwise None
-        weights: integration weight, shape (n, ). Only set when returned from Integrator.sample
-            function, otherwise None.
-        alphas: channel weights including learned correction, shape (n, channels). Only set when
-            returned from Integrator.sample function, otherwise None.
-        zero_counts: channel-wise counts of samples with zero-weights that are not included in the
-            batch, shape (channels, ). This field is ignored by most methods, as it behaves
-            does not have the batch size as its first dimension
-    """
-
-    x: torch.Tensor
-    y: torch.Tensor | None
-    q_sample: torch.Tensor
-    func_vals: torch.Tensor
-    channels: torch.Tensor | None
-    alphas_prior: torch.Tensor | None = None
-    alpha_channel_indices: torch.Tensor | None = None
-    integration_channels: torch.Tensor | None = None
-    weights: torch.Tensor | None = None
-    alphas: torch.Tensor | None = None
-    zero_counts: torch.Tensor | None = None
-
-    def __iter__(self) -> Iterable[torch.Tensor | None]:
-        """
-        Returns iterator over the fields of the class
-        """
-        return iter(astuple(self)[:-1])
-
-    def map(self, func: Callable[[torch.Tensor], torch.Tensor]) -> "SampleBatch":
-        """
-        Applies function to all fields in the batch that are not None and returns a new SampleBatch
-
-        Args:
-            func: function that is applied to all fields in the batch. Expects a tensor as argument
-                and returns a new tensor
-        Returns:
-            Transformed SampleBatch
-        """
-        return SampleBatch(*(None if field is None else func(field) for field in self))
-
-    def split(self, batch_size: int) -> Iterable["SampleBatch"]:
-        """
-        Splits up the fields into batches and yields SampleBatch objects for every batch.
-
-        Args:
-            batch_size: maximal size of the batches
-        Returns:
-            Iterator over the batches
-        """
-        for batch in zip(
-            *(
-                itertools.repeat(None) if field is None else field.split(batch_size)
-                for field in self
-            )
-        ):
-            yield SampleBatch(*batch)
-
-    @staticmethod
-    def cat(batches: Iterable["SampleBatch"]) -> "SampleBatch":
-        """
-        Concatenates multiple batches. If the field zero_counts is not None, the zero_counts of
-        all batches are added.
-
-        Args:
-            batches: Iterable over SampleBatch objects
-        Return:
-            New SamplaBatch object containing the concatenated batches
-        """
-        cat_batch = SampleBatch(
-            *(
-                None if item[0] is None else torch.cat(item, dim=0)
-                for item in zip(*batches)
-            )
-        )
-        if batches[0].zero_counts is not None:
-            cat_batch.zero_counts = torch.stack(
-                [batch.zero_counts for batch in batches], dim=1
-            ).sum(dim=1)
-        return cat_batch
-
-
-class Integrator(nn.Module):
-    """
-    Implements MadNIS training and integration logic. MadNIS integrators are torch modules, so
-    their state can easily be saved and loaded using the torch.save and torch.load methods.
-    """
-
-    def __init__(
-        self,
-        integrand: Callable[[torch.Tensor], torch.Tensor] | Integrand,
-        dims: int = 0,
-        flow: Distribution | None = None,
-        flow_kwargs: dict[str, Any] = {},
-        discrete_flow_kwargs: dict[str, Any] = {},
-        discrete_model: Literal["made", "transformer"] = "made",
-        train_channel_weights: bool = True,
-        cwnet: nn.Module | None = None,
-        cwnet_kwargs: dict[str, Any] = {},
-        loss: MultiChannelLoss | None = None,
-        optimizer: (
-            Optimizer | Callable[[Iterable[nn.Parameter]], Optimizer] | None
-        ) = None,
-        batch_size: int = 1024,
-        batch_size_per_channel: int = 0,
-        learning_rate: float = 1e-3,
-        scheduler: LRScheduler | Callable[[Optimizer], LRScheduler] | None = None,
-        uniform_channel_ratio: float = 1.0,
-        integration_history_length: int = 20,
-        drop_zero_integrands: bool = False,
-        batch_size_threshold: float = 0.5,
-        buffer_capacity: int = 0,
-        minimum_buffer_size: int = 50,
-        buffered_steps: int = 0,
-        max_stored_channel_weights: int | None = None,
-        channel_dropping_threshold: float = 0.0,
-        channel_dropping_interval: int = 100,
-        channel_grouping_mode: Literal["none", "uniform", "learned"] = "none",
-        freeze_cwnet_iteration: int | None = None,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ):
-        """
-        Args:
-            integrand: the function to be integrated. In the case of a simple single-channel
-                integration, the integrand function can directly be passed to the integrator.
-                In more complicated cases, like multi-channel integrals, use the ``Integrand`` class.
-            dims: dimension of the integration space. Only required if a simple function is given
-                as integrand.
-            flow: sampling distribution used for the integration. If None, a flow is constructed
-                using the ``Flow`` class. Otherwise, it has to be compatible with a normalizing flow,
-                i.e. have the interface defined in the ``Distribution`` class.
-            flow_kwargs: If flow is None, these keyword arguments are passed to the `Flow`
-                constructor.
-            discrete_flow_kwargs: If flow is None, these keyword arguments are passed to the
-                ``MixedFlow`` or ``DiscreteMADE`` constructor.
-            train_channel_weights: If True, construct a channel weight network and train it. Only
-                necessary if cwnet is None.
-            cwnet: network used for the trainable channel weights. If None and
-                train_channel_weights is True, the cwnet is built using the ``MLP`` class.
-            cwnet_kwargs: If cwnet is None and train_channel_weights is True, these keyword
-                arguments are passed to the ``MLP`` constructor.
-            loss: Loss function used for training. If not provided, the KL divergence is chosen in
-                the single-channel case and the stratified variance is chosen in the multi-channel
-                case.
-            optimizer: optimizer for the training. Can be an optimizer object or function that is
-                called with the model parameters as argument and returns the optimizer. If None, the
-                Adam optimizer is used.
-            batch_size: Training batch size
-            batch_size_per_channel: used to compute the batch size as a function of the number of
-                active channels, ``batch_size + n_active_channels * batch_size_per_channel``
-            learning_rate: learning rate used for the Adam optimizer
-            scheduler: learning rate scheduler for the training. Can be a learning rate scheduler
-                object or a function that gets the optimizer as argument and returns the scheduler.
-                If None, a constant learning rate is used.
-            uniform_channel_ratio: part of samples in each batch that will be distributed equally
-                between all channels, value has to be between 0 and 1.
-            integration_history_length: number of batches for which the channel-wise means and
-                variances are stored. This is used for stratified sampling during integration, and
-                during the training if uniform_channel_ratio is different from one.
-            drop_zero_integrands: If True, points with integrand zero are dropped and not used for
-                the optimization.
-            batch_size_threshold: New samples are drawn until the number of samples is at least
-                batch_size_threshold * batch_size.
-            buffer_capacity: number of samples that are stored for buffered training
-            minimum_buffer_size: minimal size of the buffer to run buffered training
-            buffered_steps: number of optimization steps on buffered samples after every online
-                training step
-            max_stored_channel_weights: number of prior channel weights that are buffered for each
-                sample. If None, all prior channel weights are saved, otherwise only those for the
-                channels with the largest contributions.
-            channel_dropping_threshold: all channels which a cumulated contribution to the
-                integrand that is smaller than this threshold are dropped
-            channel_dropping_interval: number of training steps after which channel dropping
-                is performed
-            channel_grouping_mode: If "none" all channels are treated as separate channels in the
-                loss and integration, even when they grouped together. If "uniform", the channels
-                within each group are sampled with equal probability. If "learned", a discrete
-                normalizing flow is used to sample the channel index within a group.
-            freeze_cwnet_iteration: If not None, specifies the training iteration after which the
-                channel weight network is frozen
-            device: torch device used for training and integration. If None, use default device.
-            dtype: torch dtype used for training and integration. If None, use default dtype.
-        """
-        super().__init__()
-
-        if not isinstance(integrand, Integrand):
-            integrand = Integrand(integrand, dims)
-        self.integrand = integrand
-        self.multichannel = integrand.channel_count is not None
-        discrete_dims = integrand.discrete_dims
-        input_dim = integrand.input_dim
-        if integrand.channel_grouping is None or channel_grouping_mode == "none":
-            self.integration_channel_count = integrand.channel_count
-            self.group_channels = False
-            self.group_channels_uniform = False
-        elif channel_grouping_mode == "uniform":
-            self.integration_channel_count = integrand.unique_channel_count()
-            self.group_channels = True
-            self.group_channels_uniform = True
-        elif channel_grouping_mode == "learned":
-            self.integration_channel_count = integrand.unique_channel_count()
-            self.group_channels = True
-            self.group_channels_uniform = False
-            self.channel_group_dim = (
-                0
-                if integrand.discrete_dims_position == "first"
-                else input_dim - len(discrete_dims)
-            )
-            # TODO: provide default implementation of discrete prior
-            # discrete_dims.insert(0, max(len(group.channel_indices) for group in integrand.channel_grouping.groups))
-            # input_dim += 1
-        else:
-            raise ValueError(f"Unknown channel grouping mode {channel_grouping_mode}")
-
-        if self.group_channels:
-            self.register_buffer(
-                "channel_group_sizes",
-                torch.tensor(
-                    [
-                        len(group.channel_indices)
-                        for group in integrand.channel_grouping.groups
-                    ]
-                ),
-            )
-            self.register_buffer(
-                "channel_group_remap",
-                torch.zeros(
-                    (len(self.channel_group_sizes), max(self.channel_group_sizes)),
-                    dtype=torch.int64,
-                ),
-            )
-            for group in integrand.channel_grouping.groups:
-                for i, chan_index in enumerate(group.channel_indices):
-                    self.channel_group_remap[group.group_index][i] = chan_index
-
-        if flow is None:
-            channel_remap_function = (
-                None
-                if self.group_channels and not self.group_channels_uniform
-                else self.integrand.remap_channels
-            )
-            if len(discrete_dims) == 0:
-                raise NotImplementedError("removed in MadSpace version of MadNIS")
-            elif len(discrete_dims) == input_dim:
-                if discrete_model == "made":
-                    raise NotImplementedError("removed in MadSpace version of MadNIS")
-                elif discrete_model == "transformer":
-                    raise NotImplementedError("removed in MadSpace version of MadNIS")
-                else:
-                    raise ValueError("discrete_model must be 'made' or 'transformer'")
-            else:
-                discrete_kwargs = dict(
-                    prior_prob_function=integrand.discrete_prior_prob_function,
-                    **discrete_flow_kwargs,
-                )
-                if self.multichannel:
-                    discrete_kwargs["channel_remap_function"] = channel_remap_function
-                raise NotImplementedError("removed in MadSpace version of MadNIS")
-
-        if cwnet is None and train_channel_weights and self.multichannel:
-            raise NotImplementedError("removed in MadSpace version of MadNIS")
-        if cwnet is None:
-            parameters = flow.parameters()
-        else:
-            parameters = itertools.chain(flow.parameters(), cwnet.parameters())
-        if optimizer is None:
-            self.optimizer = torch.optim.Adam(parameters, learning_rate)
-        elif isinstance(optimizer, Optimizer):
-            self.optimizer = optimizer
-        else:
-            self.optimizer = optimizer(parameters)
-        if scheduler is None or isinstance(scheduler, LRScheduler):
-            self.scheduler = scheduler
-        else:
-            self.scheduler = scheduler(self.optimizer)
-
-        self.flow = flow
-        self.cwnet = cwnet
-        self.batch_size_offset = batch_size
-        self.batch_size_per_channel = batch_size_per_channel
-        self.batch_size = batch_size + batch_size_per_channel * (
-            self.integration_channel_count or 1
-        )
-        self.uniform_channel_ratio = uniform_channel_ratio
-        self.drop_zero_integrands = drop_zero_integrands
-        self.batch_size_threshold = batch_size_threshold
-        if loss is None:
-            self.loss = stratified_variance if self.multichannel else kl_divergence
-        else:
-            self.loss = loss
-
-        self.minimum_buffer_size = minimum_buffer_size
-        self.buffered_steps = buffered_steps
-        self.max_stored_channel_weights = (
-            None
-            if max_stored_channel_weights is None
-            or integrand.channel_count is None
-            or max_stored_channel_weights >= integrand.channel_count
-            else max_stored_channel_weights
-        )
-        if buffer_capacity > 0:
-            channel_count = self.max_stored_channel_weights or integrand.channel_count
-            buffer_fields = [
-                (input_dim,),
-                None if integrand.remapped_dim is None else (integrand.remapped_dim,),
-                (),
-                (),
-                None if integrand.channel_count is None else (),
-                None if not integrand.has_channel_weight_prior else (channel_count,),
-                None if self.max_stored_channel_weights is None else (channel_count,),
-                () if self.group_channels else None,
-                None,
-                None,
-            ]
-            buffer_dtypes = [
-                None,
-                None,
-                None,
-                None,
-                torch.int64,
-                None,
-                torch.int64,
-                torch.int64,
-                None,
-                None,
-            ]
-            self.buffer = Buffer(
-                buffer_capacity, buffer_fields, persistent=False, dtypes=buffer_dtypes
-            )
-        else:
-            self.buffer = None
-        self.channel_dropping_threshold = channel_dropping_threshold
-        self.channel_dropping_interval = channel_dropping_interval
-        self.freeze_cwnet_iteration = freeze_cwnet_iteration
-        hist_shape = (self.integration_channel_count or 1,)
-        self.integration_history = Buffer(
-            integration_history_length,
-            [hist_shape, hist_shape, hist_shape],
-            dtypes=[None, None, torch.int64],
-        )
-        self.step = 0
-        self.step_type_count = 0
-        if self.multichannel:
-            self.register_buffer(
-                "active_channels_mask",
-                torch.ones((self.integration_channel_count,), dtype=torch.bool),
-            )
-        # Dummy to determine device and dtype
-        self.register_buffer("dummy", torch.zeros((1,)))
-
-        if device is not None:
-            self.to(device)
-        if dtype is not None:
-            self.to(dtype)
-
-    def _get_alphas(self, samples: SampleBatch) -> torch.Tensor:
-        """
-        Runs the channel weight network and returns the normalized channel weights, taking prior
-        channel weights and dropped channels into account.
-
-        Args:
-            samples: batch of samples
-        Returns:
-            channel weights, shape (n, channels)
-        """
-        if self.cwnet is None:
-            if samples.alphas_prior is None:
-                return samples.x.new_full(
-                    (samples.x.shape[0], self.integrand.channel_count),
-                    1 / self.integrand.channel_count,
-                )
-            return self._restore_prior(samples)
-
-        if samples.alphas_prior is None:
-            alpha_prior = samples.x.new_ones(
-                (samples.x.shape[0], self.integrand.channel_count)
-            )
-        else:
-            alpha_prior = self._restore_prior(samples)
-
-        if self.group_channels:
-            active_channels_mask = self.active_channels_mask[
-                self.integrand.remap_channels(
-                    torch.arange(alpha_prior.shape[1], device=alpha_prior.device)
-                )
-            ]
-        else:
-            active_channels_mask = self.active_channels_mask
-
-        alpha = alpha_prior * active_channels_mask
-        mask = samples.func_vals != 0
-        y = samples.x if samples.y is None else samples.y
-        alpha[mask] *= self.cwnet(y[mask]).exp()
-        ret = alpha / alpha.sum(dim=1, keepdim=True)
-        return ret
-
-    def _compute_integral(
-        self, samples: SampleBatch
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Computes normalized integrand and channel-wise means, variances and counts
-
-        Args:
-            samples: batch of samples
-        Returns:
-            A tuple containing
-
-              - normalized integrand, shape (n, )
-              - channel-wise means of the integral, shape (channels, )
-              - channel-wise variances of the integral, shape (channels, )
-              - channel-wise number of samples, shape (channels, )
-        """
-        if self.multichannel:
-            alphas = torch.gather(
-                self._get_alphas(samples), index=samples.channels[:, None], dim=1
-            )[:, 0]
-            f_true = alphas * samples.func_vals
-            f_div_q = f_true.detach() / samples.q_sample
-            channels = (
-                samples.channels
-                if samples.integration_channels is None
-                else samples.integration_channels
-            )
-            counts = torch.bincount(channels, minlength=self.integration_channel_count)
-            if samples.zero_counts is not None:
-                counts += samples.zero_counts
-            means = torch.bincount(
-                channels,
-                weights=f_div_q,
-                minlength=self.integration_channel_count,
-            ) / counts.clip(min=1)
-            variances = (
-                torch.bincount(
-                    channels,
-                    weights=(f_div_q - means[channels]).square(),
-                    minlength=self.integration_channel_count,
-                )
-                / counts
-            )
-        else:
-            f_div_q = samples.func_vals / samples.q_sample
-            f_true = samples.func_vals
-            means = f_div_q.mean(dim=0, keepdim=True)
-            counts = torch.full((1,), f_div_q.shape[0], device=means.device)
-            variances = f_div_q.var(dim=0, keepdim=True)
-        return f_true, means, variances, counts
-
-    def _optimization_step(
-        self,
-        samples: SampleBatch,
-    ) -> tuple[float, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Perform one optimization step of the networks for the given samples
-
-        Args:
-            samples: batch of samples
-        Returns:
-            A tuple containing
-
-              - value of the loss
-              - channel-wise means of the integral, shape (channels, )
-              - channel-wise variances of the integral, shape (channels, )
-              - channel-wise number of samples, shape (channels, )
-        """
-        self.optimizer.zero_grad()
-        # TODO: depending on the loss function and for drop_zero_weights=False, we can encounter
-        # zero-weight events here and it might be sufficient to evaluate the flow for events with
-        # func_val != 0. That might however give wrong results for other loss functions
-        q_test = self.flow.prob(
-            samples.x,
-            channel=(
-                samples.integration_channels
-                if self.group_channels and not self.group_channels_uniform
-                else samples.channels
-            ),
-        )
-        f_true, means, variances, counts = self._compute_integral(samples)
-        loss = self.loss(
-            f_true,
-            q_test,
-            q_sample=samples.q_sample,
-            channels=(
-                samples.channels
-                if samples.integration_channels is None
-                else samples.integration_channels
-            ),
-        )
-        if loss.isnan().item():
-            warnings.warn("nan batch: skipping optimization")
-        else:
-            loss.backward()
-            self.optimizer.step()
-        return loss.item(), means, variances, counts
-
-    def _restore_prior(self, samples: SampleBatch) -> torch.Tensor:
-        """
-        Restores the full prior channel weights if only the largest channel weights and their
-        indices were saved.
-
-        Args:
-            samples: batch of samples
-        Returns:
-            Tensor of prior channel weights with shape (n, channels)
-        """
-        if samples.alpha_channel_indices is None:
-            return samples.alphas_prior
-
-        alphas_prior_reduced = samples.alphas_prior
-        epsilon = torch.finfo(alphas_prior_reduced.dtype).eps
-
-        # strategy 1: distribute difference to 1 evenly among non-stored channels
-        # n_rest = self.integrand.channel_count - self.max_stored_channel_weights
-        # alphas_prior = torch.clamp(
-        #    (1 - alphas_prior_reduced.sum(dim=1, keepdims=True)) / n_rest,
-        #    min=epsilon,
-        # ).repeat(1, self.integrand.channel_count)
-        # alphas_prior.scatter_(1, samples.alpha_channel_indices, alphas_prior_reduced)
-        # return alphas_prior
-
-        # strategy 2: set non-stored channel alphas to epsilon, normalize again
-        alphas_prior = alphas_prior_reduced.new_full(
-            (alphas_prior_reduced.shape[0], self.integrand.channel_count), epsilon
-        )
-        alphas_prior.scatter_(1, samples.alpha_channel_indices, alphas_prior_reduced)
-        return alphas_prior / alphas_prior.sum(dim=1, keepdims=True)
-
-    def _get_channel_contributions(
-        self,
-        expect_full_history: bool,
-        channel_weight_mode: Literal["variance", "mean"],
-    ) -> torch.Tensor:
-        """
-        Uses the list of saved variances or means to compute the contribution of each channel for
-        stratified sampling.
-
-        Args:
-            expect_full_history: If True, the integration history has to be full, otherwise uniform
-                weights are returned.
-            channel_weight_mode: specifies whether the channels are weighted by their mean or
-                variance. Note that weighting by mean can lead to problems for non-positive functions
-        Returns:
-            weights for sampling the channels with shape (channels,)
-        """
-        min_len = self.integration_history.capacity if expect_full_history else 1
-        if self.integration_history.size < min_len:
-            return torch.ones(
-                self.integration_channel_count,
-                device=self.dummy.device,
-                dtype=self.dummy.dtype,
-            )
-        mean_hist, var_hist, count_hist = self.integration_history
-        contrib_hist = mean_hist.abs() if channel_weight_mode == "mean" else var_hist
-        count_hist = torch.where(
-            contrib_hist.isnan(), np.nan, count_hist.to(contrib_hist.dtype)
-        )
-        hist_weights = count_hist / count_hist.nansum(dim=0)
-        return torch.nansum(hist_weights * contrib_hist, dim=0).sqrt()
-
-    def _disable_unused_channels(self) -> int:
-        """
-        Determines channels with a total relative contribution below
-        ``channel_dropping_threshold``, disables them and removes them from the buffer.
-
-        Returns:
-            Number of channels that were disabled
-        """
-        if (
-            not self.multichannel
-            or self.channel_dropping_threshold == 0.0
-            or (self.step + 1) % self.channel_dropping_interval != 0
-        ):
-            return 0
-
-        mean_hist, _, count_hist = self.integration_history
-        count_hist = count_hist.to(mean_hist.dtype)
-        mean_hist = torch.nan_to_num(mean_hist)
-        hist_weights = count_hist / count_hist.sum(dim=0)
-        channel_integrals = torch.nansum(hist_weights * mean_hist, dim=0)
-        channel_rel_integrals = channel_integrals / channel_integrals.sum()
-        cri_sort, cri_argsort = torch.sort(channel_rel_integrals)
-        n_irrelevant = torch.count_nonzero(
-            cri_sort.cumsum(dim=0) < self.channel_dropping_threshold
-        )
-        n_disabled = torch.count_nonzero(
-            self.active_channels_mask[cri_argsort[:n_irrelevant]]
-        )
-        self.active_channels_mask[cri_argsort[:n_irrelevant]] = False
-        self.integrand.update_active_channels_mask(self.active_channels_mask)
-        self.batch_size = (
-            self.batch_size_offset
-            + torch.count_nonzero(self.active_channels_mask).item()
-            * self.batch_size_per_channel
-        )
-        if self.buffer is not None:
-
-            def filter_func(batch):
-                samples = SampleBatch(*batch)
-                channels = (
-                    samples.channels
-                    if samples.integration_channels is None
-                    else samples.integration_channels
-                )
-                return self.active_channels_mask[channels]
-
-            self.buffer.filter(filter_func)
-        return n_disabled
-
-    def _store_samples(self, samples: SampleBatch):
-        """
-        Stores the generated samples and probabilites for reuse during buffered training. If
-        ``max_stored_channel_weights`` is set, the largest channel weights are determined and only
-        those and their weights are stored.
-
-        Args:
-            samples: Object containing a batch of samples
-        """
-        if self.buffer is None:
-            return
-
-        if (
-            self.max_stored_channel_weights is not None
-            and self.integrand.has_channel_weight_prior
-        ):
-            # ensure that the alpha for the channel that the sample was generated with
-            # is always stored
-            alphas_prior_mod = torch.scatter(
-                samples.alphas_prior,
-                dim=1,
-                index=samples.channels[:, None],
-                src=torch.tensor(
-                    [[2.0]], device=self.dummy.device, dtype=self.dummy.dtype
-                ).expand(*samples.alphas_prior.shape),
-            )
-            largest_alphas, alpha_indices = torch.sort(
-                alphas_prior_mod, descending=True, dim=1
-            )
-            largest_alphas[:, 0] = torch.gather(
-                samples.alphas_prior, dim=1, index=samples.channels[:, None]
-            )[:, 0]
-            samples.alphas_prior = largest_alphas[
-                :, : self.max_stored_channel_weights
-            ].clone()
-            samples.alpha_channel_indices = alpha_indices[
-                :, : self.max_stored_channel_weights
-            ].clone()
-
-        self.buffer.store(*samples)
-
-    def _get_channels(
-        self,
-        n: int,
-        channel_weights: torch.Tensor,
-        uniform_channel_ratio: float,
-        return_counts: bool = False,
-    ) -> torch.Tensor:
-        """
-        Create a tensor of channel indices or number of samples per channel in two steps:
-        1. Split up n * uniform_channel_ratio equally among all the channels
-        2. Sample the rest of the events from the distribution given by channel_weights
-           after correcting for the uniformly distributed samples
-        This allows stratified sampling by variance weighting while ensuring stable training
-        because there are events in every channel.
-        Args:
-            n: Number of samples as scalar integer tensor
-            channel_weights: Weights of the channels (not normalized) with shape (channels,)
-            uniform_channel_ratio: Number between 0.0 and 1.0 to determine the ratio of samples
-                that will be distributed uniformly first
-            return_counts: If True, return number of samples per channels, otherwise the channel
-                indices
-        Returns:
-            If return_counts is True, Tensor with number of samples per channel, shape (channels,).
-            Otherwise, Tensor of channel numbers with shape (n,)
-        """
-        assert channel_weights.shape == (self.integration_channel_count,)
-        n_active_channels = torch.count_nonzero(self.active_channels_mask).item()
-        uniform_per_channel = int(
-            np.ceil(n * uniform_channel_ratio / n_active_channels)
-        )
-        n_per_channel = torch.full(
-            (self.integration_channel_count,),
-            uniform_per_channel,
-            device=self.dummy.device,
-        )
-        n_per_channel[~self.active_channels_mask] = 0
-
-        n_weighted = max(n - n_per_channel.sum(), 0)
-        if n_weighted > 0:
-            normed_weights = (
-                channel_weights / channel_weights[self.active_channels_mask].sum()
-            )
-            normed_weights[~self.active_channels_mask] = 0.0
-            probs = torch.clamp(
-                normed_weights - uniform_channel_ratio / n_active_channels, min=0
-            )
-            n_per_channel += torch.ceil(probs * n_weighted / probs.sum()).int()
-
-        remove_chan = 0
-        while n_per_channel.sum() > n:
-            if n_per_channel[remove_chan] > 0:
-                n_per_channel[remove_chan] -= 1
-            remove_chan = (remove_chan + 1) % self.integration_channel_count
-        assert n_per_channel.sum() == n
-
-        if return_counts:
-            return n_per_channel
-
-        return torch.cat(
-            [
-                torch.full((npc,), i, device=self.dummy.device)
-                for i, npc in enumerate(n_per_channel)
-            ]
-        )
-
-    def _get_samples(
-        self,
-        n: int,
-        uniform_channel_ratio: float = 0.0,
-        train: bool = False,
-        channel_weight_mode: Literal["variance", "mean"] = "variance",
-        channel: int | None = None,
-    ) -> SampleBatch:
-        """
-        Draws samples from the flow and evaluates the integrand
-
-        Args:
-            n: number of samples
-            uniform_channel_ratio: Number between 0.0 and 1.0 to determine the ratio of samples
-                that will be distributed uniformly first
-            train: If True, the function is used in training mode, i.e. samples where the integrand
-                is zero will be removed if drop_zero_integrands is True
-            channel_weight_mode: specifies whether the channels are weighted by their mean or
-                variance. Note that weighting by mean can lead to problems for non-positive functions
-            channel: if different from None, samples are only generated for this channel
-        Returns:
-            Object containing a batch of samples
-        """
-        if channel is None:
-            batch_channels = (
-                self._get_channels(
-                    n,
-                    self._get_channel_contributions(train, channel_weight_mode),
-                    uniform_channel_ratio,
-                )
-                if self.multichannel
-                else None
-            )
-        else:
-            batch_channels = torch.full((n,), channel, device=self.dummy.device)
-
-        batches_out = []
-        current_batch_size = 0
-        while True:
-            integration_channels = None
-            weight_factor = None
-            if self.integrand.function_includes_sampling:
-                integration_channels = batch_channels
-                x, prob, weight, y, alphas_prior, channels = self.integrand.function(
-                    batch_channels
-                )
-            else:
-                if self.group_channels and self.group_channels_uniform:
-                    group_sizes = self.channel_group_sizes[batch_channels]
-                    chan_in_group = (
-                        torch.rand(
-                            (n,), device=self.dummy.device, dtype=self.dummy.dtype
-                        )
-                        * group_sizes
-                    ).long()
-                    weight_factor = group_sizes
-                    integration_channels = batch_channels
-                    channels = self.channel_group_remap[batch_channels, chan_in_group]
-                else:
-                    channels = batch_channels
-
-                with torch.no_grad():
-                    x, prob = self.flow.sample(
-                        n,
-                        channel=channels,
-                        return_prob=True,
-                        device=self.dummy.device,
-                        dtype=self.dummy.dtype,
-                    )
-                weight, y, alphas_prior = self.integrand(x, channels)
-
-                if self.group_channels and not self.group_channels_uniform:
-                    chan_in_group = x[:, self.channel_group_dim].long()
-                    integration_channels = batch_channels
-                    channels = self.channel_group_remap[
-                        integration_channels, chan_in_group
-                    ]
-
-            if weight_factor is not None:
-                weight *= weight_factor
-            batch = SampleBatch(
-                x,
-                y,
-                prob,
-                weight,
-                channels,
-                alphas_prior,
-                integration_channels=integration_channels,
-            )
-
-            if not train:
-                current_batch_size += batch.x.shape[0]
-            elif self.drop_zero_integrands:
-                mask = weight != 0.0
-                batch = batch.map(lambda t: t[mask])
-                if self.multichannel:
-                    # batch.zero_counts = torch.bincount(
-                    #    (
-                    #        channels[~mask]
-                    #        if integration_channels is None
-                    #        else integration_channels[~mask]
-                    #    ),
-                    #    minlength=self.integration_channel_count,
-                    # )
-                    c = (
-                        channels[mask]
-                        if integration_channels is None
-                        else integration_channels[mask]
-                    )
-                    batch.q_sample /= (
-                        torch.bincount(
-                            c,
-                            minlength=self.integration_channel_count,
-                        )
-                        / torch.bincount(
-                            c,
-                            minlength=self.integration_channel_count,
-                        )
-                    )[c]
-                else:
-                    # batch.zero_counts = torch.full(
-                    #    (1,), torch.count_nonzero(~mask), device=x.device
-                    # )
-                    batch.q_sample /= mask.double().mean()
-                current_batch_size += batch.x.shape[0]
-            else:
-                current_batch_size += weight.count_nonzero()
-
-            # mask = ~(weight.isnan() | x.isnan().any(dim=1)) & mask
-            batches_out.append(batch)
-
-            # check this condition at the end such that the sampling runs at least once
-            if current_batch_size > self.batch_size_threshold * n:
-                break
-
-        return SampleBatch.cat(batches_out)
-
-    def train_step(self) -> TrainingStatus:
-        """
-        Performs a single training step
-
-        Returns:
-            Training status
-        """
-
-        if self.step == self.freeze_cwnet_iteration and self.cwnet is not None:
-            for param in self.cwnet.parameters():
-                param.requires_grad = False
-
-        if self.step_type_count == 0:
-            buffered = False
-            samples = self._get_samples(
-                self.batch_size, self.uniform_channel_ratio, train=True
-            )
-            loss, means, variances, counts = self._optimization_step(samples)
-            self._store_samples(samples)
-            self.integration_history.store(means[None], variances[None], counts[None])
-
-            if self.buffered_steps != 0 and self.buffer.size > self.minimum_buffer_size:
-                self.step_type_count += 1
-        else:
-            buffered = True
-            samples = SampleBatch(*self.buffer.sample(self.batch_size))
-            loss, _, _, _ = self._optimization_step(samples)
-            self.step_type_count = (self.step_type_count + 1) % (
-                self.buffered_steps + 1
-            )
-
-        dropped_channels = self._disable_unused_channels()
-        status = TrainingStatus(
-            step=self.step,
-            loss=loss,
-            buffered=buffered,
-            learning_rate=(
-                None if self.scheduler is None else self.scheduler.get_last_lr()[0]
-            ),
-            dropped_channels=dropped_channels,
-        )
-
-        if self.scheduler is not None:
-            self.scheduler.step()
-        self.step += 1
-        return status
-
-    def train(
-        self,
-        steps: int,
-        callback: Callable[[TrainingStatus], None] | None = None,
-        capture_keyboard_interrupt: bool = False,
-    ):
-        """
-        Performs multiple training steps
-
-        Args:
-            steps: number of training steps
-            callback: function that is called after each training step with the training status
-                as argument
-            capture_keyboard_interrupt: If True, a keyboard interrupt does not raise an exception.
-                Instead, the current training step is finished and the training is aborted
-                afterwards.
-        """
-        interrupted = False
-        if capture_keyboard_interrupt:
-
-            def handler(sig, frame):
-                nonlocal interrupted
-                interrupted = True
-
-            old_handler = signal.signal(signal.SIGINT, handler)
-
-        try:
-            for _ in range(steps):
-                status = self.train_step()
-                if callback is not None:
-                    callback(status)
-                if interrupted:
-                    break
-        finally:
-            if capture_keyboard_interrupt:
-                signal.signal(signal.SIGINT, old_handler)
diff --git a/madspace/madspace/madnis/interface.py b/madspace/madspace/madnis/interface.py
deleted file mode 100644
index b36c82f770..0000000000
--- a/madspace/madspace/madnis/interface.py
+++ /dev/null
@@ -1,177 +0,0 @@
-from collections.abc import Callable
-
-import torch
-import torch.nn as nn
-
-from .. import _madspace_py_loader as ms
-from ..torch import FunctionModule
-from .channel_grouping import ChannelGrouping
-from .distribution import Distribution
-from .integrand import Integrand
-
-
-class IntegrandDistribution(nn.Module, Distribution):
-    def __init__(
-        self,
-        channels: list[ms.Integrand],
-        channel_remap_function: Callable[[torch.Tensor], torch.Tensor],
-        context: ms.Context,
-    ):
-        super().__init__()
-        self.channel_count = len(channels)
-        self.channels = channels
-        self.context = context
-        self.channel_remap_function = channel_remap_function
-        self.latent_dims, self.latent_float = channels[0].latent_dims()
-        self.integrand_prob = None
-        self.update_channel_mask(torch.ones(self.channel_count, dtype=torch.bool))
-
-    def update_channel_mask(self, mask: torch.Tensor) -> None:
-        self.channel_mask = mask
-        multi_prob = ms.MultiChannelFunction(
-            [
-                ms.IntegrandProbability(chan)
-                for chan, active in zip(self.channels, mask)
-                if active
-            ]
-        )
-        func = multi_prob.function()
-        if self.integrand_prob is None:
-            self.integrand_prob = FunctionModule(func, self.context)
-        else:
-            self.integrand_prob.runtime = ms.FunctionRuntime(func, self.context)
-
-    def sample(
-        self,
-        n: int,
-        c: torch.Tensor | None = None,
-        channel: torch.Tensor | list[int] | int | None = None,
-        return_log_prob: bool = False,
-        return_prob: bool = False,
-        device: torch.device | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
-        raise NotImplementedError(
-            "IntegrandDistribution does not support sampling directly. "
-            "Use the underlying ms.Integrand object instead."
-        )
-
-    def prob(
-        self,
-        x: torch.Tensor,
-        c: torch.Tensor | None = None,
-        channel: torch.Tensor | list[int] | int | None = None,
-    ) -> torch.Tensor:
-        channel_perm = None
-        if isinstance(channel, torch.Tensor):
-            channel = self.channel_remap_function(channel)
-            channel_perm = torch.argsort(channel)
-            x = x[channel_perm]
-            channel = channel.bincount(minlength=self.channel_count).to(torch.int32)
-        elif channel is None:
-            channel = torch.tensor([len(x)], dtype=torch.int32)
-        else:
-            raise NotImplementedError("channel argument type not supported")
-        channel = channel[self.channel_mask]
-
-        prob_args = [
-            xi if is_float else xi[:, 0].to(torch.int32)
-            for xi, is_float in zip(x.split(self.latent_dims, dim=1), self.latent_float)
-        ]
-        prob = self.integrand_prob(*prob_args, channel.cpu())
-        if channel_perm is None:
-            return prob
-        else:
-            channel_perm_inv = torch.argsort(channel_perm)
-            return prob[channel_perm_inv]
-
-
-class IntegrandFunction:
-    def __init__(self, channels: list[ms.Integrand], context: ms.Context):
-        self.channel_count = len(channels)
-        self.channels = channels
-        self.context = context
-        self.update_channel_mask(torch.ones(self.channel_count, dtype=torch.bool))
-
-    def update_channel_mask(self, mask: torch.Tensor) -> None:
-        self.channel_mask = mask.cpu()
-        multi_integrand = ms.MultiChannelIntegrand(
-            [chan for chan, active in zip(self.channels, mask) if active]
-        )
-        self.multi_runtime = ms.FunctionRuntime(
-            multi_integrand.function(), self.context
-        )
-
-    def __call__(self, channels: torch.Tensor) -> tuple[torch.Tensor, ...]:
-        channel_perm = torch.argsort(channels)
-        channels = channels.bincount(minlength=self.channel_count).cpu().to(torch.int32)
-        channels = channels[self.channel_mask]
-        (
-            weight,
-            latent,
-            prob,
-            chan_index,
-            alphas_prior,
-            y,
-            *rest,
-        ) = self.multi_runtime(channels)
-
-        x_parts = [latent, *rest]
-        x = torch.cat(
-            [xi.double().reshape(latent.shape[0], -1) for xi in x_parts], dim=1
-        )
-        channel_perm_inv = torch.argsort(channel_perm)
-        return (
-            x[channel_perm_inv],
-            prob[channel_perm_inv],
-            weight[channel_perm_inv],
-            y[channel_perm_inv],
-            alphas_prior[channel_perm_inv],
-            chan_index[channel_perm_inv].to(torch.int64),
-        )
-
-
-def build_madnis_integrand(
-    channels: list[ms.Integrand],
-    cwnet: ms.ChannelWeightNetwork | None = None,
-    channel_grouping: ChannelGrouping | None = None,
-    context: ms.Context = ms.default_context(),
-) -> tuple[Integrand, Distribution, nn.Module | None]:
-    device = torch.device("cpu" if context.device() == ms.cpu_device() else "cuda:0")
-    if channel_grouping is None:
-        remap_channels = lambda channels: channels
-        group_indices = torch.arange(len(channels))
-    else:
-        channel_id_map = torch.tensor(
-            [channel.group.group_index for channel in channel_grouping.channels],
-            device=device,
-        )
-        remap_channels = lambda channels: channel_id_map[channels]
-        group_indices = torch.tensor(
-            [group.target_index for group in channel_grouping.groups], device=device
-        )
-
-    integrand_function = IntegrandFunction(channels, context)
-    flow = IntegrandDistribution(channels, remap_channels, context)
-
-    def update_mask(mask: torch.Tensor) -> None:
-        context.get_global(cwnet.mask_name()).torch()[0, :] = mask.double()
-        group_mask = mask[group_indices]
-        if torch.any(group_mask.cpu() != integrand_function.channel_mask):
-            integrand_function.update_channel_mask(group_mask)
-            flow.update_channel_mask(group_mask)
-
-    integrand = Integrand(
-        function=integrand_function,
-        input_dim=sum(channels[0].latent_dims()[0]),
-        channel_count=len(channel_grouping.channels),
-        remapped_dim=cwnet.preprocessing().output_dim(),
-        has_channel_weight_prior=cwnet is not None,
-        channel_grouping=channel_grouping,
-        function_includes_sampling=True,
-        update_active_channels_mask=update_mask,
-    )
-    cwnet_module = (
-        None if cwnet is None else FunctionModule(cwnet.mlp().function(), context)
-    )
-    return integrand, flow, cwnet_module
diff --git a/madspace/madspace/madnis/losses.py b/madspace/madspace/madnis/losses.py
deleted file mode 100644
index 9daaaa3c2b..0000000000
--- a/madspace/madspace/madnis/losses.py
+++ /dev/null
@@ -1,224 +0,0 @@
-"""Implementation of multi-channel loss functions"""
-
-from collections.abc import Callable
-from functools import wraps
-
-import torch
-
-
-def dtype_epsilon(tensor: torch.Tensor) -> float:
-    return torch.finfo(tensor.dtype).eps
-
-
-def softclip(x: torch.Tensor, threshold: torch.Tensor = 30.0):
-    return threshold * torch.arcsinh(x / threshold)
-
-
-SingleChannelLoss = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
-MultiChannelLoss = Callable[
-    [torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None], torch.Tensor
-]
-
-
-def multi_channel_loss(loss: SingleChannelLoss) -> MultiChannelLoss:
-    """
-    Turns a single-channel loss function into a multi-channel loss function by evaluating it for
-    each channel separately and then adding them weighted by TODO weighted by what?
-
-    Args:
-        loss: single-channel loss function, that expects the integrand value, test probability and
-            sampling probability as arguments
-    Returns:
-        multi-channel loss function, that expects the integrand value, test probability and,
-        optionally, sampling probability and channel indices as arguments.
-    """
-
-    # TODO: this unfortunately does not yield the correct signature (with the extra channels argument),
-    # so it does not show up in the documentation
-    @wraps(loss)
-    def wrapped_multi(
-        f_true: torch.Tensor,
-        q_test: torch.Tensor,
-        q_sample: torch.Tensor | None = None,
-        channels: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        if q_sample is None:
-            q_sample = q_test
-        if channels is None:
-            return loss(f_true, q_test, q_sample)
-
-        loss_tot = 0
-        for channel in channels.unique():
-            mask = channels == channel
-            fi, qti, qsi = f_true[mask], q_test[mask], q_sample[mask]
-            ni = mask.count_nonzero()
-            # loss_tot += ni / q_sample.shape[0] * loss(fi, qti, qsi) if ni > 0 else 0.0
-            loss_tot += loss(fi, qti, qsi) if ni > 0 else 0.0
-        return loss_tot
-
-    return wrapped_multi
-
-
-def stratified_variance(
-    f_true: torch.Tensor,
-    q_test: torch.Tensor,
-    q_sample: torch.Tensor | None = None,
-    channels: torch.Tensor | None = None,
-):
-    """
-    Computes the stratified variance as introduced in [2311.01548] for two given sets of
-    probabilities, ``f_true`` and ``q_test``. It uses importance sampling with a sampling
-    probability specified by ``q_sample``.
-
-    Args:
-        f_true: normalized integrand values
-        q_test: estimated function/probability
-        q_sample: sampling probability
-        channels: channel indices or None in the single-channel case
-    Returns:
-        computed stratified variance
-    """
-    if q_sample is None:
-        q_sample = q_test
-    if channels is None:
-        abs_integral = torch.mean(f_true.detach().abs() / q_sample)
-        return _variance(f_true, q_test, q_sample) / abs_integral.square()
-
-    stddev_sum = 0
-    abs_integral = 0
-    for i in channels.unique():
-        mask = channels == i
-        fi, qti, qsi = f_true[mask], q_test[mask], q_sample[mask]
-        stddev_sum += torch.sqrt(_variance(fi, qti, qsi) + dtype_epsilon(f_true))
-        abs_integral += torch.mean(fi.detach().abs() / qsi)
-    return (stddev_sum / abs_integral) ** 2
-
-
-def stratified_variance_softclip(
-    f_true: torch.Tensor,
-    q_test: torch.Tensor,
-    q_sample: torch.Tensor | None = None,
-    channels: torch.Tensor | None = None,
-    threshold: torch.Tensor = 30.0,
-):
-    """
-    Computes the stratified variance as introduced in [2311.01548] for two given sets of
-    probabilities, ``f_true`` and ``q_test``. It uses importance sampling with a sampling
-    probability specified by ``q_sample``. A soft clipping function is applied to the
-    sample weights.
-
-    Args:
-        f_true: normalized integrand values
-        q_test: estimated function/probability
-        q_sample: sampling probability
-        channels: channel indices or None in the single-channel case
-        threshold: approximate point of transition between linear and logarithmic behavior
-    Returns:
-        computed stratified variance
-    """
-    if q_sample is None:
-        q_sample = q_test
-    if channels is None:
-        norm = torch.mean(f_true.detach().abs() / q_sample)
-        f_true = softclip(f_true / q_sample / norm) * q_sample * norm
-        abs_integral = torch.mean(f_true.detach().abs() / q_sample)
-        return _variance(f_true, q_test, q_sample) / abs_integral.square()
-
-    stddev_sum = 0
-    abs_integral = 0
-    for i in channels.unique():
-        mask = channels == i
-        fi, qti, qsi = f_true[mask], q_test[mask], q_sample[mask]
-        norm = torch.mean(fi.detach().abs() / qsi)
-        fi = softclip(fi / qsi / norm) * qsi * norm
-        stddev_sum += torch.sqrt(_variance(fi, qti, qsi) + dtype_epsilon(f_true))
-        abs_integral += torch.mean(fi.detach().abs() / qsi)
-    return (stddev_sum / abs_integral) ** 2
-
-
-@multi_channel_loss
-def variance(
-    f_true: torch.Tensor, q_test: torch.Tensor, q_sample: torch.Tensor
-) -> torch.Tensor:
-    abs_integral = torch.mean(f_true.detach().abs() / q_sample) + dtype_epsilon(f_true)
-    return _variance(f_true, q_test, q_sample) / abs_integral.square()
-
-
-def _variance(
-    f_true: torch.Tensor,
-    q_test: torch.Tensor,
-    q_sample: torch.Tensor,
-) -> torch.Tensor:
-    """
-    Computes the variance for two given sets of probabilities, ``f_true`` and ``q_test``. It uses
-    importance sampling with a sampling probability specified by ``q_sample``.
-
-    Args:
-        f_true: normalized integrand values
-        q_test: estimated function/probability
-        q_sample: sampling probability
-    Returns:
-        computed variance
-    """
-    ratio = q_test / q_sample
-    mean = torch.mean(f_true / q_sample)
-    sq = (f_true / q_test - mean) ** 2
-    return (
-        torch.mean(sq * ratio)
-        if len(f_true) > 0
-        else torch.tensor(0.0, device=f_true.device, dtype=f_true.dtype)
-    )
-
-
-@multi_channel_loss
-def kl_divergence(
-    f_true: torch.Tensor, q_test: torch.Tensor, q_sample: torch.Tensor
-) -> torch.Tensor:
-    """
-    Computes the Kullback-Leibler divergence for two given sets of probabilities, ``f_true`` and
-    ``q_test``. It uses importance sampling, i.e. the estimator is divided by an additional factor
-    of ``q_sample``.
-
-    Args:
-        f_true: normalized integrand values
-        q_test: estimated function/probability
-        q_sample: sampling probability
-        channels: channel indices or None in the single-channel case
-    Returns:
-        computed KL divergence
-    """
-    f_true = f_true.detach().abs()
-    f_true /= torch.mean(f_true / q_sample)
-    log_q = torch.log(q_test)
-    log_f = torch.log(f_true + dtype_epsilon(f_true))
-    return torch.mean(f_true / q_sample * (log_f - log_q))
-
-
-@multi_channel_loss
-def kl_divergence_softclip(
-    f_true: torch.Tensor,
-    q_test: torch.Tensor,
-    q_sample: torch.Tensor,
-    threshold: torch.Tensor = 30.0,
-) -> torch.Tensor:
-    """
-    Computes the Kullback-Leibler divergence for two given sets of probabilities, ``f_true`` and
-    ``q_test``. It uses importance sampling, i.e. the estimator is divided by an additional factor
-    of ``q_sample``. A soft clipping function is applied to the sample weights.
-
-    Args:
-        f_true: normalized integrand values
-        q_test: estimated function/probability
-        q_sample: sampling probability
-        channels: channel indices or None in the single-channel case
-        threshold: approximate point of transition between linear and logarithmic behavior
-    Returns:
-        computed KL divergence
-    """
-    f_true = f_true.detach().abs()
-    weight = f_true / q_sample
-    weight /= weight.abs().mean()
-    clipped_weight = softclip(weight, threshold)
-    log_q = torch.log(q_test)
-    log_f = torch.log(clipped_weight * q_sample + dtype_epsilon(f_true))
-    return torch.mean(clipped_weight * (log_f - log_q))
diff --git a/madspace/src/compgraphs/instruction_set_mixin.inc b/madspace/src/compgraphs/instruction_set_mixin.inc
index a2dbe6cd88..083f13abb2 100644
--- a/madspace/src/compgraphs/instruction_set_mixin.inc
+++ b/madspace/src/compgraphs/instruction_set_mixin.inc
@@ -36,140 +36,141 @@ InstructionOwner instructions[] {
     mi("div", 18, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
     mi("reduce_sum", 19, true, {{DataType::dt_float, false, {std::monostate{}, "n"}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
     mi("reduce_sum_vector", 20, true, {{DataType::dt_float, false, {std::monostate{}, "n", "m"}, false}}, {{DataType::dt_float, false, {std::monostate{}, "m"}, false}}),
-    mi("batch_reduce_mean", 21, true, {{DataType::dt_float, false, {}, false}}, {{DataType::dt_float, true, {}, false}}),
-    mi("batch_reduce_mean_keepdim", 22, true, {{DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("reduce_product", 23, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("sqrt", 24, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("square", 25, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("min", 26, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("max", 27, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_sqrt_s", 28, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("obs_e", 29, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_px", 30, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_py", 31, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_pz", 32, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_mass", 33, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_pt", 34, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_p_mag", 35, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_phi", 36, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_theta", 37, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_y", 38, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_y_abs", 39, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_eta", 40, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_eta_abs", 41, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_delta_eta", 42, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}, {DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_delta_phi", 43, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}, {DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("obs_delta_r", 44, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}, {DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("boost_beam", 45, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}),
-    mi("boost_beam_inverse", 46, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}),
-    mi("com_p_in", 47, true, {{DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}),
-    mi("r_to_x1x2", 48, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("x1x2_to_r", 49, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("diff_cross_section", 50, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("two_body_decay_com", 51, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("two_body_decay_com_inverse", 52, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("two_body_decay", 53, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("two_body_decay_inverse", 54, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("two_to_two_particle_scattering_com", 55, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("two_to_two_particle_scattering_com_inverse", 56, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("two_to_two_particle_scattering", 57, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("two_to_two_particle_scattering_inverse", 58, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("two_to_three_particle_scattering", 59, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("two_to_three_particle_scattering_inverse", 60, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("double_t_scattering", 61, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("double_t_scattering_inverse", 62, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("three_body_decay_com", 63, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("three_body_decay_com_inverse", 64, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("three_body_decay", 65, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("three_body_decay_inverse", 66, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("t_inv_min_max", 67, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("t_inv_value_and_min_max", 68, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("t_inv_min_max_cut", 69, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("t_inv_value_and_min_max_cut", 70, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("t1_inv_min_max_doublet", 71, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("t1_inv_value_and_min_max_doublet", 72, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("t2_inv_min_max_doublet", 73, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("t2_inv_value_and_min_max_doublet", 74, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("s23_min_max", 75, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("s23_value_and_min_max", 76, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("s23_min_max_cut", 77, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("s23_value_and_min_max_cut", 78, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("invariants_from_momenta", 79, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {"m", "n"}, false}}, {{DataType::dt_float, false, {"m"}, false}}),
-    mi("sde2_channel_weights", 80, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}}, {{DataType::dt_float, false, {"c"}, false}}),
-    mi("subchannel_weights", 81, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, true, {"g"}, false}}, {{DataType::dt_float, false, {"c"}, false}}),
-    mi("apply_subchannel_weights", 82, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"d"}, false}, {DataType::dt_int, false, {"d"}, false}}, {{DataType::dt_float, false, {"d"}, false}}),
-    mi("pt_eta_phi_x", 83, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n+2"}, false}}),
-    mi("mirror_momenta", 84, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}),
-    mi("momenta_to_x1x2", 85, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("uniform_invariant", 86, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("uniform_invariant_inverse", 87, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("breit_wigner_invariant", 88, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("breit_wigner_invariant_inverse", 89, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("stable_invariant", 90, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("stable_invariant_inverse", 91, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("stable_invariant_nu", 92, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("stable_invariant_nu_inverse", 93, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("fast_rambo_massless", 94, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("fast_rambo_massless_inverse", 95, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("fast_rambo_massless_com", 96, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("fast_rambo_massive", 97, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("fast_rambo_massive_inverse", 98, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("fast_rambo_massive_com", 99, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("cut_unphysical", 100, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("cut_one", 101, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("cut_all", 102, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("cut_any", 103, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("scale_transverse_energy", 104, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("scale_transverse_mass", 105, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("scale_half_transverse_mass", 106, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("scale_partonic_energy", 107, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("chili_forward", 108, true, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("chili_inverse", 109, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}}),
-    InstructionOwner(new MatrixElementInstruction(110, true)),
-    mi("collect_channel_weights", 111, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}, {DataType::dt_int, true, {"c"}, true}}, {{DataType::dt_float, false, {"c"}, false}}),
-    mi("interpolate_pdf", 112, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {"n"}, false}, {DataType::dt_float, true, {"a"}, false}, {DataType::dt_float, true, {"b"}, false}, {DataType::dt_float, true, {16, "c", "d"}, false}}, {{DataType::dt_float, false, {"n"}, false}}),
-    mi("interpolate_alpha_s", 113, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {"b+1"}, false}, {DataType::dt_float, true, {4, "b"}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("matmul", 114, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"m", "n"}, false}, {DataType::dt_float, true, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}),
-    mi("relu", 115, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("leaky_relu", 116, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("elu", 117, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("gelu", 118, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("sigmoid", 119, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("softplus", 120, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    InstructionOwner(new RqsReshapeInstruction(121, true)),
-    mi("rqs_find_bin", 122, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b+1"}, false}}, {{DataType::dt_float, false, {"n", 6}, false}}),
-    mi("rqs_forward", 123, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}),
-    mi("rqs_inverse", 124, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}),
-    mi("softmax", 125, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
-    mi("softmax_prior", 126, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}),
-    mi("sample_discrete", 127, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("sample_discrete_inverse", 128, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("sample_discrete_probs", 129, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("sample_discrete_probs_inverse", 130, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
-    mi("discrete_histogram", 131, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, true, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}}),
-    mi("permute_momenta", 132, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, true, {"m", "n"}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}),
-    mi("gather", 133, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("gather_int", 134, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}),
-    mi("gather_vector", 135, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n", "m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}),
-    mi("select_int", 136, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}),
-    mi("select", 137, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}),
-    mi("select_vector", 138, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}),
-    mi("argsort", 139, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}),
-    mi("quantile", 140, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, true, {}, false}}),
-    mi("one_hot", 141, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}),
-    mi("madnis_abs_weight", 142, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("madnis_softclip", 143, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("madnis_variance", 144, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("madnis_single_channel_variance", 145, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
-    mi("madnis_multi_channel_variance", 146, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}),
-    InstructionOwner(new NonzeroInstruction(147, true)),
-    InstructionOwner(new BatchGatherInstruction(148, true)),
-    InstructionOwner(new BatchScatterInstruction(149, true)),
-    InstructionOwner(new RandomInstruction(150, true)),
-    InstructionOwner(new RandomIntInstruction(151, true)),
-    InstructionOwner(new UnweightInstruction(152, true)),
-    mi("vegas_forward", 153, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}),
-    mi("vegas_inverse", 154, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}),
-    mi("vegas_histogram", 155, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"n", "b"}, false}, {DataType::dt_int, true, {"n", "b"}, false}}),
-    mi("histogram", 156, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"b+2"}, false}, {DataType::dt_float, true, {"b+2"}, false}}),
+    mi("batch_reduce_sum", 21, true, {{DataType::dt_float, false, {}, false}}, {{DataType::dt_float, true, {}, false}}),
+    mi("batch_reduce_mean", 22, true, {{DataType::dt_float, false, {}, false}}, {{DataType::dt_float, true, {}, false}}),
+    mi("batch_reduce_mean_keepdim", 23, true, {{DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("reduce_product", 24, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("sqrt", 25, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("square", 26, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("min", 27, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("max", 28, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_sqrt_s", 29, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("obs_e", 30, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_px", 31, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_py", 32, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_pz", 33, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_mass", 34, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_pt", 35, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_p_mag", 36, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_phi", 37, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_theta", 38, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_y", 39, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_y_abs", 40, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_eta", 41, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_eta_abs", 42, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_delta_eta", 43, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}, {DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_delta_phi", 44, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}, {DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("obs_delta_r", 45, true, {{DataType::dt_float, false, {std::monostate{}, 4}, false}, {DataType::dt_float, false, {std::monostate{}, 4}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("boost_beam", 46, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}),
+    mi("boost_beam_inverse", 47, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}),
+    mi("com_p_in", 48, true, {{DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}),
+    mi("r_to_x1x2", 49, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("x1x2_to_r", 50, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("diff_cross_section", 51, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("two_body_decay_com", 52, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("two_body_decay_com_inverse", 53, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("two_body_decay", 54, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("two_body_decay_inverse", 55, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("two_to_two_particle_scattering_com", 56, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("two_to_two_particle_scattering_com_inverse", 57, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("two_to_two_particle_scattering", 58, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("two_to_two_particle_scattering_inverse", 59, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("two_to_three_particle_scattering", 60, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("two_to_three_particle_scattering_inverse", 61, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("double_t_scattering", 62, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("double_t_scattering_inverse", 63, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("three_body_decay_com", 64, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("three_body_decay_com_inverse", 65, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("three_body_decay", 66, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("three_body_decay_inverse", 67, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("t_inv_min_max", 68, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("t_inv_value_and_min_max", 69, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("t_inv_min_max_cut", 70, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("t_inv_value_and_min_max_cut", 71, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("t1_inv_min_max_doublet", 72, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("t1_inv_value_and_min_max_doublet", 73, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("t2_inv_min_max_doublet", 74, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("t2_inv_value_and_min_max_doublet", 75, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("s23_min_max", 76, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("s23_value_and_min_max", 77, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("s23_min_max_cut", 78, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("s23_value_and_min_max_cut", 79, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("invariants_from_momenta", 80, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {"m", "n"}, false}}, {{DataType::dt_float, false, {"m"}, false}}),
+    mi("sde2_channel_weights", 81, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}}, {{DataType::dt_float, false, {"c"}, false}}),
+    mi("subchannel_weights", 82, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, true, {"g"}, false}}, {{DataType::dt_float, false, {"c"}, false}}),
+    mi("apply_subchannel_weights", 83, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"d"}, false}, {DataType::dt_int, false, {"d"}, false}}, {{DataType::dt_float, false, {"d"}, false}}),
+    mi("pt_eta_phi_x", 84, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n+2"}, false}}),
+    mi("mirror_momenta", 85, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}),
+    mi("momenta_to_x1x2", 86, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("uniform_invariant", 87, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("uniform_invariant_inverse", 88, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("breit_wigner_invariant", 89, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("breit_wigner_invariant_inverse", 90, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("stable_invariant", 91, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("stable_invariant_inverse", 92, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("stable_invariant_nu", 93, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("stable_invariant_nu_inverse", 94, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("fast_rambo_massless", 95, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("fast_rambo_massless_inverse", 96, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("fast_rambo_massless_com", 97, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("fast_rambo_massive", 98, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("fast_rambo_massive_inverse", 99, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("fast_rambo_massive_com", 100, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("cut_unphysical", 101, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("cut_one", 102, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("cut_all", 103, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("cut_any", 104, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("scale_transverse_energy", 105, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("scale_transverse_mass", 106, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("scale_half_transverse_mass", 107, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("scale_partonic_energy", 108, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("chili_forward", 109, true, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("chili_inverse", 110, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}}),
+    InstructionOwner(new MatrixElementInstruction(111, true)),
+    mi("collect_channel_weights", 112, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}, {DataType::dt_int, true, {"c"}, true}}, {{DataType::dt_float, false, {"c"}, false}}),
+    mi("interpolate_pdf", 113, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {"n"}, false}, {DataType::dt_float, true, {"a"}, false}, {DataType::dt_float, true, {"b"}, false}, {DataType::dt_float, true, {16, "c", "d"}, false}}, {{DataType::dt_float, false, {"n"}, false}}),
+    mi("interpolate_alpha_s", 114, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {"b+1"}, false}, {DataType::dt_float, true, {4, "b"}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("matmul", 115, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"m", "n"}, false}, {DataType::dt_float, true, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}),
+    mi("relu", 116, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("leaky_relu", 117, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("elu", 118, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("gelu", 119, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("sigmoid", 120, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("softplus", 121, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    InstructionOwner(new RqsReshapeInstruction(122, true)),
+    mi("rqs_find_bin", 123, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b+1"}, false}}, {{DataType::dt_float, false, {"n", 6}, false}}),
+    mi("rqs_forward", 124, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}),
+    mi("rqs_inverse", 125, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}),
+    mi("softmax", 126, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}),
+    mi("softmax_prior", 127, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}),
+    mi("sample_discrete", 128, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("sample_discrete_inverse", 129, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("sample_discrete_probs", 130, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("sample_discrete_probs_inverse", 131, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}),
+    mi("discrete_histogram", 132, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, true, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}}),
+    mi("permute_momenta", 133, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, true, {"m", "n"}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}),
+    mi("gather", 134, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("gather_int", 135, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}),
+    mi("gather_vector", 136, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n", "m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}),
+    mi("select_int", 137, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}),
+    mi("select", 138, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}),
+    mi("select_vector", 139, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}),
+    mi("argsort", 140, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}),
+    mi("quantile", 141, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, true, {}, false}}),
+    mi("one_hot", 142, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}),
+    mi("madnis_abs_weight", 143, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("madnis_softclip", 144, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("madnis_variance", 145, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("madnis_single_channel_variance", 146, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}),
+    mi("madnis_multi_channel_variance", 147, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}),
+    InstructionOwner(new NonzeroInstruction(148, true)),
+    InstructionOwner(new BatchGatherInstruction(149, true)),
+    InstructionOwner(new BatchScatterInstruction(150, true)),
+    InstructionOwner(new RandomInstruction(151, true)),
+    InstructionOwner(new RandomIntInstruction(152, true)),
+    InstructionOwner(new UnweightInstruction(153, true)),
+    mi("vegas_forward", 154, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}),
+    mi("vegas_inverse", 155, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}),
+    mi("vegas_histogram", 156, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"n", "b"}, false}, {DataType::dt_int, true, {"n", "b"}, false}}),
+    mi("histogram", 157, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"b+2"}, false}, {DataType::dt_float, true, {"b+2"}, false}}),
 };
diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp
index 02ed5a5be7..9f0fb2bde9 100644
--- a/madspace/src/cpu/runtime.cpp
+++ b/madspace/src/cpu/runtime.cpp
@@ -408,11 +408,12 @@ void op_batch_scatter(
 }
 
 template <typename D>
-void batch_reduce_mean_impl(
+void batch_reduce_sum_mean_impl(
     const CpuRuntime::Instruction& instruction,
     TensorVec& locals,
     const D& device,
-    bool keepdim
+    bool keepdim,
+    bool is_mean
 ) {
     auto& input = locals[instruction.input_indices[0]];
     auto& output = locals[instruction.output_indices[0]];
@@ -422,30 +423,39 @@ void batch_reduce_mean_impl(
 
     auto input_view_flat = input.flat_view<double, 1>(0);
     auto output_view_flat = output.flat_view<double, 1>(0);
-    device.submit([keepdim, input_view_flat, output_view_flat, batch_size]() mutable {
-        TensorView<double, 1> input_view(input_view_flat);
-        TensorView<double, 1> output_view(output_view_flat);
-        double sum = 0.;
-        for (std::size_t i = 0; i < batch_size; ++i) {
-            sum += input_view[i];
-        }
-        if (keepdim) {
+    device.submit(
+        [keepdim, is_mean, input_view_flat, output_view_flat, batch_size]() mutable {
+            TensorView<double, 1> input_view(input_view_flat);
+            TensorView<double, 1> output_view(output_view_flat);
+            double sum = 0.;
             for (std::size_t i = 0; i < batch_size; ++i) {
-                output_view[i] = sum / batch_size;
+                sum += input_view[i];
+            }
+            if (keepdim) {
+                if (is_mean) {
+                    for (std::size_t i = 0; i < batch_size; ++i) {
+                        output_view[i] = sum / batch_size;
+                    }
+                } else {
+                    for (std::size_t i = 0; i < batch_size; ++i) {
+                        output_view[i] = sum;
+                    }
+                }
+            } else {
+                output_view[0] = is_mean ? sum / batch_size : sum;
             }
-        } else {
-            output_view[0] = sum / batch_size;
         }
-    });
+    );
 }
 
 template <typename D>
-void batch_reduce_mean_backward_impl(
+void batch_reduce_sum_mean_backward_impl(
     const CpuRuntime::Instruction& instruction,
     TensorVec& locals,
     TensorVec& local_grads,
     const D& device,
-    bool keepdim
+    bool keepdim,
+    bool is_mean
 ) {
     auto& input = locals[instruction.input_indices[0]];
     auto& input_grad = local_grads[instruction.input_indices[0]];
@@ -461,7 +471,11 @@ void batch_reduce_mean_backward_impl(
     auto output_grad_view_flat = output_grad.flat_view<double, 1>(0);
     std::size_t batch_size = input.size(0);
     device.submit(
-        [keepdim, input_grad_view_flat, output_grad_view_flat, batch_size]() mutable {
+        [keepdim,
+         is_mean,
+         input_grad_view_flat,
+         output_grad_view_flat,
+         batch_size]() mutable {
             TensorView<double, 1> input_grad_view(input_grad_view_flat);
             TensorView<double, 1> output_grad_view(output_grad_view_flat);
             double grad = 0.0;
@@ -469,9 +483,11 @@ void batch_reduce_mean_backward_impl(
                 for (std::size_t i = 0; i < batch_size; ++i) {
                     grad += output_grad_view[i];
                 }
-                grad /= batch_size;
             } else {
-                grad = output_grad_view[0] / batch_size;
+                grad = output_grad_view[0];
+            }
+            if (is_mean) {
+                grad /= batch_size;
             }
             for (std::size_t i = 0; i < batch_size; ++i) {
                 input_grad_view[i] += grad;
@@ -480,11 +496,30 @@ void batch_reduce_mean_backward_impl(
     );
 }
 
+template <typename D>
+void op_batch_reduce_sum(
+    const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device
+) {
+    batch_reduce_sum_mean_impl(instruction, locals, device, false, false);
+}
+
+template <typename D>
+void backward_op_batch_reduce_sum(
+    const CpuRuntime::Instruction& instruction,
+    TensorVec& locals,
+    TensorVec& local_grads,
+    const D& device
+) {
+    batch_reduce_sum_mean_backward_impl(
+        instruction, locals, local_grads, device, false, false
+    );
+}
+
 template <typename D>
 void op_batch_reduce_mean(
     const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device
 ) {
-    batch_reduce_mean_impl(instruction, locals, device, false);
+    batch_reduce_sum_mean_impl(instruction, locals, device, false, true);
 }
 
 template <typename D>
@@ -494,14 +529,16 @@ void backward_op_batch_reduce_mean(
     TensorVec& local_grads,
     const D& device
 ) {
-    batch_reduce_mean_backward_impl(instruction, locals, local_grads, device, false);
+    batch_reduce_sum_mean_backward_impl(
+        instruction, locals, local_grads, device, false, true
+    );
 }
 
 template <typename D>
 void op_batch_reduce_mean_keepdim(
     const CpuRuntime::Instruction& instruction, TensorVec& locals, const D& device
 ) {
-    batch_reduce_mean_impl(instruction, locals, device, true);
+    batch_reduce_sum_mean_impl(instruction, locals, device, true, true);
 }
 
 template <typename D>
@@ -511,7 +548,9 @@ void backward_op_batch_reduce_mean_keepdim(
     TensorVec& local_grads,
     const D& device
 ) {
-    batch_reduce_mean_backward_impl(instruction, locals, local_grads, device, true);
+    batch_reduce_sum_mean_backward_impl(
+        instruction, locals, local_grads, device, true, true
+    );
 }
 
 template <typename D>
diff --git a/madspace/src/cpu/runtime_backward_mixin.inc b/madspace/src/cpu/runtime_backward_mixin.inc
index 8a75bf4987..f7cb036ead 100644
--- a/madspace/src/cpu/runtime_backward_mixin.inc
+++ b/madspace/src/cpu/runtime_backward_mixin.inc
@@ -38,80 +38,83 @@ case 18:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_div<CpuTypes>, backward_kernel_div<SimdTypes>, 3, 2, DeviceType>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
 case 21:
-    backward_op_batch_reduce_mean(instr, locals, local_grads, device);
+    backward_op_batch_reduce_sum(instr, locals, local_grads, device);
     break;
 case 22:
-    backward_op_batch_reduce_mean_keepdim(instr, locals, local_grads, device);
+    backward_op_batch_reduce_mean(instr, locals, local_grads, device);
     break;
 case 23:
-    backward_batch_foreach<tensor_foreach<backward_kernel_reduce_product<CpuTypes>, backward_kernel_reduce_product<SimdTypes>, 2, 1, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
+    backward_op_batch_reduce_mean_keepdim(instr, locals, local_grads, device);
     break;
 case 24:
-    backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_sqrt<CpuTypes>, backward_kernel_sqrt<SimdTypes>, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
+    backward_batch_foreach<tensor_foreach<backward_kernel_reduce_product<CpuTypes>, backward_kernel_reduce_product<SimdTypes>, 2, 1, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
 case 25:
+    backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_sqrt<CpuTypes>, backward_kernel_sqrt<SimdTypes>, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
+    break;
+case 26:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_square<CpuTypes>, backward_kernel_square<SimdTypes>, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 114:
+case 115:
     backward_op_matmul(instr, locals, local_grads, device);
     break;
-case 115:
+case 116:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_relu<CpuTypes>, backward_kernel_relu<SimdTypes>, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 116:
+case 117:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_leaky_relu<CpuTypes>, backward_kernel_leaky_relu<SimdTypes>, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 117:
+case 118:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_elu<CpuTypes>, backward_kernel_elu<SimdTypes>, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 118:
+case 119:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_gelu<CpuTypes>, backward_kernel_gelu<SimdTypes>, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 119:
+case 120:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_sigmoid<CpuTypes>, backward_kernel_sigmoid<SimdTypes>, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device);
     break;
-case 120:
+case 121:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_softplus<CpuTypes>, backward_kernel_softplus<SimdTypes>, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 121:
+case 122:
     backward_op_rqs_reshape(instr, locals, local_grads, device);
     break;
-case 122:
+case 123:
     backward_batch_foreach<tensor_foreach<backward_kernel_rqs_find_bin<CpuTypes>, backward_kernel_rqs_find_bin<SimdTypes>, 5, 4, 2, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device);
     break;
-case 123:
+case 124:
     backward_batch_foreach<tensor_foreach<backward_kernel_rqs_forward<CpuTypes>, backward_kernel_rqs_forward<SimdTypes>, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
-case 124:
+case 125:
     backward_batch_foreach<tensor_foreach<backward_kernel_rqs_inverse<CpuTypes>, backward_kernel_rqs_inverse<SimdTypes>, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
-case 125:
+case 126:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_softmax<CpuTypes>, backward_kernel_softmax<SimdTypes>, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device);
     break;
-case 126:
+case 127:
     backward_batch_foreach<tensor_foreach<backward_kernel_softmax_prior<CpuTypes>, backward_kernel_softmax_prior<SimdTypes>, 2, 2, 1, DeviceType>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device);
     break;
-case 130:
+case 131:
     backward_batch_foreach<tensor_foreach<backward_kernel_sample_discrete_probs_inverse<CpuTypes>, backward_kernel_sample_discrete_probs_inverse<SimdTypes>, 4, 2, 1, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
-case 133:
+case 134:
     backward_batch_foreach<tensor_foreach<backward_kernel_gather<CpuTypes>, backward_kernel_gather<SimdTypes>, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 137:
+case 138:
     backward_batch_foreach<tensor_foreach<backward_kernel_select<CpuTypes>, backward_kernel_select<SimdTypes>, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device);
     break;
-case 142:
+case 143:
     backward_batch_foreach<tensor_foreach<backward_kernel_madnis_abs_weight<CpuTypes>, backward_kernel_madnis_abs_weight<SimdTypes>, 3, 2, 1, DeviceType>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
-case 143:
+case 144:
     backward_batch_foreach<tensor_foreach<backward_kernel_madnis_softclip<CpuTypes>, backward_kernel_madnis_softclip<SimdTypes>, 5, 4, 1, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device);
     break;
-case 144:
+case 145:
     backward_batch_foreach<tensor_foreach<backward_kernel_madnis_variance<CpuTypes>, backward_kernel_madnis_variance<SimdTypes>, 5, 4, 1, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device);
     break;
-case 145:
+case 146:
     backward_batch_foreach<tensor_foreach<backward_kernel_madnis_single_channel_variance<CpuTypes>, backward_kernel_madnis_single_channel_variance<SimdTypes>, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device);
     break;
-case 146:
+case 147:
     backward_batch_foreach<tensor_foreach<backward_kernel_madnis_multi_channel_variance<CpuTypes>, backward_kernel_madnis_multi_channel_variance<SimdTypes>, 3, 2, 1, DeviceType>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
diff --git a/madspace/src/cpu/runtime_mixin.inc b/madspace/src/cpu/runtime_mixin.inc
index 64a3974be4..24ba7d2f30 100644
--- a/madspace/src/cpu/runtime_mixin.inc
+++ b/madspace/src/cpu/runtime_mixin.inc
@@ -65,410 +65,413 @@ case 20:
     batch_foreach<tensor_foreach_dynamic<kernel_reduce_sum_vector<CpuTypes>, kernel_reduce_sum_vector<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 21:
-    op_batch_reduce_mean(instr, locals, device);
+    op_batch_reduce_sum(instr, locals, device);
     break;
 case 22:
-    op_batch_reduce_mean_keepdim(instr, locals, device);
+    op_batch_reduce_mean(instr, locals, device);
     break;
 case 23:
-    batch_foreach<tensor_foreach<kernel_reduce_product<CpuTypes>, kernel_reduce_product<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    op_batch_reduce_mean_keepdim(instr, locals, device);
     break;
 case 24:
-    batch_foreach<tensor_foreach_dynamic<kernel_sqrt<CpuTypes>, kernel_sqrt<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_reduce_product<CpuTypes>, kernel_reduce_product<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 25:
-    batch_foreach<tensor_foreach_dynamic<kernel_square<CpuTypes>, kernel_square<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_sqrt<CpuTypes>, kernel_sqrt<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 26:
-    batch_foreach<tensor_foreach_dynamic<kernel_min<CpuTypes>, kernel_min<SimdTypes>, 2, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_square<CpuTypes>, kernel_square<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 27:
-    batch_foreach<tensor_foreach_dynamic<kernel_max<CpuTypes>, kernel_max<SimdTypes>, 2, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_min<CpuTypes>, kernel_min<SimdTypes>, 2, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 28:
-    batch_foreach<tensor_foreach<kernel_obs_sqrt_s<CpuTypes>, kernel_obs_sqrt_s<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_max<CpuTypes>, kernel_max<SimdTypes>, 2, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 29:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_e<CpuTypes>, kernel_obs_e<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_obs_sqrt_s<CpuTypes>, kernel_obs_sqrt_s<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 30:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_px<CpuTypes>, kernel_obs_px<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_e<CpuTypes>, kernel_obs_e<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 31:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_py<CpuTypes>, kernel_obs_py<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_px<CpuTypes>, kernel_obs_px<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 32:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_pz<CpuTypes>, kernel_obs_pz<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_py<CpuTypes>, kernel_obs_py<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 33:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_mass<CpuTypes>, kernel_obs_mass<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_pz<CpuTypes>, kernel_obs_pz<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 34:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_pt<CpuTypes>, kernel_obs_pt<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_mass<CpuTypes>, kernel_obs_mass<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 35:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_p_mag<CpuTypes>, kernel_obs_p_mag<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_pt<CpuTypes>, kernel_obs_pt<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 36:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_phi<CpuTypes>, kernel_obs_phi<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_p_mag<CpuTypes>, kernel_obs_p_mag<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 37:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_theta<CpuTypes>, kernel_obs_theta<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_phi<CpuTypes>, kernel_obs_phi<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 38:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_y<CpuTypes>, kernel_obs_y<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_theta<CpuTypes>, kernel_obs_theta<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 39:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_y_abs<CpuTypes>, kernel_obs_y_abs<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_y<CpuTypes>, kernel_obs_y<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 40:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_eta<CpuTypes>, kernel_obs_eta<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_y_abs<CpuTypes>, kernel_obs_y_abs<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 41:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_eta_abs<CpuTypes>, kernel_obs_eta_abs<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_eta<CpuTypes>, kernel_obs_eta<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 42:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_eta<CpuTypes>, kernel_obs_delta_eta<SimdTypes>, 2, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_eta_abs<CpuTypes>, kernel_obs_eta_abs<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 43:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_phi<CpuTypes>, kernel_obs_delta_phi<SimdTypes>, 2, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_eta<CpuTypes>, kernel_obs_delta_eta<SimdTypes>, 2, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 44:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_r<CpuTypes>, kernel_obs_delta_r<SimdTypes>, 2, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_phi<CpuTypes>, kernel_obs_delta_phi<SimdTypes>, 2, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 45:
-    batch_foreach<tensor_foreach<kernel_boost_beam<CpuTypes>, kernel_boost_beam<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_r<CpuTypes>, kernel_obs_delta_r<SimdTypes>, 2, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 46:
-    batch_foreach<tensor_foreach<kernel_boost_beam_inverse<CpuTypes>, kernel_boost_beam_inverse<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_boost_beam<CpuTypes>, kernel_boost_beam<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
     break;
 case 47:
-    batch_foreach<tensor_foreach<kernel_com_p_in<CpuTypes>, kernel_com_p_in<SimdTypes>, 1, 2, 1, DeviceType>, 1, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_boost_beam_inverse<CpuTypes>, kernel_boost_beam_inverse<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
     break;
 case 48:
-    batch_foreach<tensor_foreach<kernel_r_to_x1x2<CpuTypes>, kernel_r_to_x1x2<SimdTypes>, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_com_p_in<CpuTypes>, kernel_com_p_in<SimdTypes>, 1, 2, 1, DeviceType>, 1, 2>(instr, locals, device);
     break;
 case 49:
-    batch_foreach<tensor_foreach<kernel_x1x2_to_r<CpuTypes>, kernel_x1x2_to_r<SimdTypes>, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_r_to_x1x2<CpuTypes>, kernel_r_to_x1x2<SimdTypes>, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device);
     break;
 case 50:
-    batch_foreach<tensor_foreach<kernel_diff_cross_section<CpuTypes>, kernel_diff_cross_section<SimdTypes>, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_x1x2_to_r<CpuTypes>, kernel_x1x2_to_r<SimdTypes>, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device);
     break;
 case 51:
-    batch_foreach<tensor_foreach<kernel_two_body_decay_com<CpuTypes>, kernel_two_body_decay_com<SimdTypes>, 5, 3, 1, DeviceType>, 5, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_diff_cross_section<CpuTypes>, kernel_diff_cross_section<SimdTypes>, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device);
     break;
 case 52:
-    batch_foreach<tensor_foreach<kernel_two_body_decay_com_inverse<CpuTypes>, kernel_two_body_decay_com_inverse<SimdTypes>, 2, 6, 1, DeviceType>, 2, 6>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_body_decay_com<CpuTypes>, kernel_two_body_decay_com<SimdTypes>, 5, 3, 1, DeviceType>, 5, 3>(instr, locals, device);
     break;
 case 53:
-    batch_foreach<tensor_foreach<kernel_two_body_decay<CpuTypes>, kernel_two_body_decay<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_body_decay_com_inverse<CpuTypes>, kernel_two_body_decay_com_inverse<SimdTypes>, 2, 6, 1, DeviceType>, 2, 6>(instr, locals, device);
     break;
 case 54:
-    batch_foreach<tensor_foreach<kernel_two_body_decay_inverse<CpuTypes>, kernel_two_body_decay_inverse<SimdTypes>, 2, 7, 1, DeviceType>, 2, 7>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_body_decay<CpuTypes>, kernel_two_body_decay<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
     break;
 case 55:
-    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_com<CpuTypes>, kernel_two_to_two_particle_scattering_com<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_body_decay_inverse<CpuTypes>, kernel_two_body_decay_inverse<SimdTypes>, 2, 7, 1, DeviceType>, 2, 7>(instr, locals, device);
     break;
 case 56:
-    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_com_inverse<CpuTypes>, kernel_two_to_two_particle_scattering_com_inverse<SimdTypes>, 4, 4, 1, DeviceType>, 4, 4>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_com<CpuTypes>, kernel_two_to_two_particle_scattering_com<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
     break;
 case 57:
-    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering<CpuTypes>, kernel_two_to_two_particle_scattering<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_com_inverse<CpuTypes>, kernel_two_to_two_particle_scattering_com_inverse<SimdTypes>, 4, 4, 1, DeviceType>, 4, 4>(instr, locals, device);
     break;
 case 58:
-    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_inverse<CpuTypes>, kernel_two_to_two_particle_scattering_inverse<SimdTypes>, 4, 4, 1, DeviceType>, 4, 4>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering<CpuTypes>, kernel_two_to_two_particle_scattering<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
     break;
 case 59:
-    batch_foreach<tensor_foreach<kernel_two_to_three_particle_scattering<CpuTypes>, kernel_two_to_three_particle_scattering<SimdTypes>, 8, 3, 1, DeviceType>, 8, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_inverse<CpuTypes>, kernel_two_to_two_particle_scattering_inverse<SimdTypes>, 4, 4, 1, DeviceType>, 4, 4>(instr, locals, device);
     break;
 case 60:
-    batch_foreach<tensor_foreach<kernel_two_to_three_particle_scattering_inverse<CpuTypes>, kernel_two_to_three_particle_scattering_inverse<SimdTypes>, 7, 4, 1, DeviceType>, 7, 4>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_three_particle_scattering<CpuTypes>, kernel_two_to_three_particle_scattering<SimdTypes>, 8, 3, 1, DeviceType>, 8, 3>(instr, locals, device);
     break;
 case 61:
-    batch_foreach<tensor_foreach<kernel_double_t_scattering<CpuTypes>, kernel_double_t_scattering<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_three_particle_scattering_inverse<CpuTypes>, kernel_two_to_three_particle_scattering_inverse<SimdTypes>, 7, 4, 1, DeviceType>, 7, 4>(instr, locals, device);
     break;
 case 62:
-    batch_foreach<tensor_foreach<kernel_double_t_scattering_inverse<CpuTypes>, kernel_double_t_scattering_inverse<SimdTypes>, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_double_t_scattering<CpuTypes>, kernel_double_t_scattering<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
     break;
 case 63:
-    batch_foreach<tensor_foreach<kernel_three_body_decay_com<CpuTypes>, kernel_three_body_decay_com<SimdTypes>, 9, 4, 1, DeviceType>, 9, 4>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_double_t_scattering_inverse<CpuTypes>, kernel_double_t_scattering_inverse<SimdTypes>, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device);
     break;
 case 64:
-    batch_foreach<tensor_foreach<kernel_three_body_decay_com_inverse<CpuTypes>, kernel_three_body_decay_com_inverse<SimdTypes>, 3, 10, 1, DeviceType>, 3, 10>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_three_body_decay_com<CpuTypes>, kernel_three_body_decay_com<SimdTypes>, 9, 4, 1, DeviceType>, 9, 4>(instr, locals, device);
     break;
 case 65:
-    batch_foreach<tensor_foreach<kernel_three_body_decay<CpuTypes>, kernel_three_body_decay<SimdTypes>, 10, 4, 1, DeviceType>, 10, 4>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_three_body_decay_com_inverse<CpuTypes>, kernel_three_body_decay_com_inverse<SimdTypes>, 3, 10, 1, DeviceType>, 3, 10>(instr, locals, device);
     break;
 case 66:
-    batch_foreach<tensor_foreach<kernel_three_body_decay_inverse<CpuTypes>, kernel_three_body_decay_inverse<SimdTypes>, 3, 11, 1, DeviceType>, 3, 11>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_three_body_decay<CpuTypes>, kernel_three_body_decay<SimdTypes>, 10, 4, 1, DeviceType>, 10, 4>(instr, locals, device);
     break;
 case 67:
-    batch_foreach<tensor_foreach<kernel_t_inv_min_max<CpuTypes>, kernel_t_inv_min_max<SimdTypes>, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_three_body_decay_inverse<CpuTypes>, kernel_three_body_decay_inverse<SimdTypes>, 3, 11, 1, DeviceType>, 3, 11>(instr, locals, device);
     break;
 case 68:
-    batch_foreach<tensor_foreach<kernel_t_inv_value_and_min_max<CpuTypes>, kernel_t_inv_value_and_min_max<SimdTypes>, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t_inv_min_max<CpuTypes>, kernel_t_inv_min_max<SimdTypes>, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device);
     break;
 case 69:
-    batch_foreach<tensor_foreach<kernel_t_inv_min_max_cut<CpuTypes>, kernel_t_inv_min_max_cut<SimdTypes>, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t_inv_value_and_min_max<CpuTypes>, kernel_t_inv_value_and_min_max<SimdTypes>, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device);
     break;
 case 70:
-    batch_foreach<tensor_foreach<kernel_t_inv_value_and_min_max_cut<CpuTypes>, kernel_t_inv_value_and_min_max_cut<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t_inv_min_max_cut<CpuTypes>, kernel_t_inv_min_max_cut<SimdTypes>, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device);
     break;
 case 71:
-    batch_foreach<tensor_foreach<kernel_t1_inv_min_max_doublet<CpuTypes>, kernel_t1_inv_min_max_doublet<SimdTypes>, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t_inv_value_and_min_max_cut<CpuTypes>, kernel_t_inv_value_and_min_max_cut<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
     break;
 case 72:
-    batch_foreach<tensor_foreach<kernel_t1_inv_value_and_min_max_doublet<CpuTypes>, kernel_t1_inv_value_and_min_max_doublet<SimdTypes>, 7, 3, 1, DeviceType>, 7, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t1_inv_min_max_doublet<CpuTypes>, kernel_t1_inv_min_max_doublet<SimdTypes>, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device);
     break;
 case 73:
-    batch_foreach<tensor_foreach<kernel_t2_inv_min_max_doublet<CpuTypes>, kernel_t2_inv_min_max_doublet<SimdTypes>, 7, 2, 1, DeviceType>, 7, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t1_inv_value_and_min_max_doublet<CpuTypes>, kernel_t1_inv_value_and_min_max_doublet<SimdTypes>, 7, 3, 1, DeviceType>, 7, 3>(instr, locals, device);
     break;
 case 74:
-    batch_foreach<tensor_foreach<kernel_t2_inv_value_and_min_max_doublet<CpuTypes>, kernel_t2_inv_value_and_min_max_doublet<SimdTypes>, 8, 3, 1, DeviceType>, 8, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t2_inv_min_max_doublet<CpuTypes>, kernel_t2_inv_min_max_doublet<SimdTypes>, 7, 2, 1, DeviceType>, 7, 2>(instr, locals, device);
     break;
 case 75:
-    batch_foreach<tensor_foreach<kernel_s23_min_max<CpuTypes>, kernel_s23_min_max<SimdTypes>, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t2_inv_value_and_min_max_doublet<CpuTypes>, kernel_t2_inv_value_and_min_max_doublet<SimdTypes>, 8, 3, 1, DeviceType>, 8, 3>(instr, locals, device);
     break;
 case 76:
-    batch_foreach<tensor_foreach<kernel_s23_value_and_min_max<CpuTypes>, kernel_s23_value_and_min_max<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_s23_min_max<CpuTypes>, kernel_s23_min_max<SimdTypes>, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device);
     break;
 case 77:
-    batch_foreach<tensor_foreach<kernel_s23_min_max_cut<CpuTypes>, kernel_s23_min_max_cut<SimdTypes>, 10, 2, 1, DeviceType>, 10, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_s23_value_and_min_max<CpuTypes>, kernel_s23_value_and_min_max<SimdTypes>, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device);
     break;
 case 78:
-    batch_foreach<tensor_foreach<kernel_s23_value_and_min_max_cut<CpuTypes>, kernel_s23_value_and_min_max_cut<SimdTypes>, 10, 3, 1, DeviceType>, 10, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_s23_min_max_cut<CpuTypes>, kernel_s23_min_max_cut<SimdTypes>, 10, 2, 1, DeviceType>, 10, 2>(instr, locals, device);
     break;
 case 79:
-    batch_foreach<tensor_foreach<kernel_invariants_from_momenta<CpuTypes>, kernel_invariants_from_momenta<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_s23_value_and_min_max_cut<CpuTypes>, kernel_s23_value_and_min_max_cut<SimdTypes>, 10, 3, 1, DeviceType>, 10, 3>(instr, locals, device);
     break;
 case 80:
-    batch_foreach<tensor_foreach<kernel_sde2_channel_weights<CpuTypes>, kernel_sde2_channel_weights<SimdTypes>, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_invariants_from_momenta<CpuTypes>, kernel_invariants_from_momenta<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 81:
-    batch_foreach<tensor_foreach<kernel_subchannel_weights<CpuTypes>, kernel_subchannel_weights<SimdTypes>, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_sde2_channel_weights<CpuTypes>, kernel_sde2_channel_weights<SimdTypes>, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device);
     break;
 case 82:
-    batch_foreach<tensor_foreach<kernel_apply_subchannel_weights<CpuTypes>, kernel_apply_subchannel_weights<SimdTypes>, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_subchannel_weights<CpuTypes>, kernel_subchannel_weights<SimdTypes>, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device);
     break;
 case 83:
-    batch_foreach<tensor_foreach<kernel_pt_eta_phi_x<CpuTypes>, kernel_pt_eta_phi_x<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_apply_subchannel_weights<CpuTypes>, kernel_apply_subchannel_weights<SimdTypes>, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device);
     break;
 case 84:
-    batch_foreach<tensor_foreach<kernel_mirror_momenta<CpuTypes>, kernel_mirror_momenta<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_pt_eta_phi_x<CpuTypes>, kernel_pt_eta_phi_x<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
     break;
 case 85:
-    batch_foreach<tensor_foreach<kernel_momenta_to_x1x2<CpuTypes>, kernel_momenta_to_x1x2<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_mirror_momenta<CpuTypes>, kernel_mirror_momenta<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 86:
-    batch_foreach<tensor_foreach<kernel_uniform_invariant<CpuTypes>, kernel_uniform_invariant<SimdTypes>, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_momenta_to_x1x2<CpuTypes>, kernel_momenta_to_x1x2<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
     break;
 case 87:
-    batch_foreach<tensor_foreach<kernel_uniform_invariant_inverse<CpuTypes>, kernel_uniform_invariant_inverse<SimdTypes>, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_uniform_invariant<CpuTypes>, kernel_uniform_invariant<SimdTypes>, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device);
     break;
 case 88:
-    batch_foreach<tensor_foreach<kernel_breit_wigner_invariant<CpuTypes>, kernel_breit_wigner_invariant<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_uniform_invariant_inverse<CpuTypes>, kernel_uniform_invariant_inverse<SimdTypes>, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device);
     break;
 case 89:
-    batch_foreach<tensor_foreach<kernel_breit_wigner_invariant_inverse<CpuTypes>, kernel_breit_wigner_invariant_inverse<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_breit_wigner_invariant<CpuTypes>, kernel_breit_wigner_invariant<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
     break;
 case 90:
-    batch_foreach<tensor_foreach<kernel_stable_invariant<CpuTypes>, kernel_stable_invariant<SimdTypes>, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_breit_wigner_invariant_inverse<CpuTypes>, kernel_breit_wigner_invariant_inverse<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
     break;
 case 91:
-    batch_foreach<tensor_foreach<kernel_stable_invariant_inverse<CpuTypes>, kernel_stable_invariant_inverse<SimdTypes>, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_stable_invariant<CpuTypes>, kernel_stable_invariant<SimdTypes>, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device);
     break;
 case 92:
-    batch_foreach<tensor_foreach<kernel_stable_invariant_nu<CpuTypes>, kernel_stable_invariant_nu<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_stable_invariant_inverse<CpuTypes>, kernel_stable_invariant_inverse<SimdTypes>, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device);
     break;
 case 93:
-    batch_foreach<tensor_foreach<kernel_stable_invariant_nu_inverse<CpuTypes>, kernel_stable_invariant_nu_inverse<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_stable_invariant_nu<CpuTypes>, kernel_stable_invariant_nu<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
     break;
 case 94:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massless<CpuTypes>, kernel_fast_rambo_massless<SimdTypes>, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_stable_invariant_nu_inverse<CpuTypes>, kernel_stable_invariant_nu_inverse<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
     break;
 case 95:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massless_inverse<CpuTypes>, kernel_fast_rambo_massless_inverse<SimdTypes>, 2, 3, 1, DeviceType>, 2, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massless<CpuTypes>, kernel_fast_rambo_massless<SimdTypes>, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device);
     break;
 case 96:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massless_com<CpuTypes>, kernel_fast_rambo_massless_com<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massless_inverse<CpuTypes>, kernel_fast_rambo_massless_inverse<SimdTypes>, 2, 3, 1, DeviceType>, 2, 3>(instr, locals, device);
     break;
 case 97:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massive<CpuTypes>, kernel_fast_rambo_massive<SimdTypes>, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massless_com<CpuTypes>, kernel_fast_rambo_massless_com<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
     break;
 case 98:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massive_inverse<CpuTypes>, kernel_fast_rambo_massive_inverse<SimdTypes>, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massive<CpuTypes>, kernel_fast_rambo_massive<SimdTypes>, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device);
     break;
 case 99:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massive_com<CpuTypes>, kernel_fast_rambo_massive_com<SimdTypes>, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massive_inverse<CpuTypes>, kernel_fast_rambo_massive_inverse<SimdTypes>, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device);
     break;
 case 100:
-    batch_foreach<tensor_foreach<kernel_cut_unphysical<CpuTypes>, kernel_cut_unphysical<SimdTypes>, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massive_com<CpuTypes>, kernel_fast_rambo_massive_com<SimdTypes>, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device);
     break;
 case 101:
-    batch_foreach<tensor_foreach<kernel_cut_one<CpuTypes>, kernel_cut_one<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_cut_unphysical<CpuTypes>, kernel_cut_unphysical<SimdTypes>, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device);
     break;
 case 102:
-    batch_foreach<tensor_foreach<kernel_cut_all<CpuTypes>, kernel_cut_all<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_cut_one<CpuTypes>, kernel_cut_one<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
     break;
 case 103:
-    batch_foreach<tensor_foreach<kernel_cut_any<CpuTypes>, kernel_cut_any<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_cut_all<CpuTypes>, kernel_cut_all<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
     break;
 case 104:
-    batch_foreach<tensor_foreach<kernel_scale_transverse_energy<CpuTypes>, kernel_scale_transverse_energy<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_cut_any<CpuTypes>, kernel_cut_any<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
     break;
 case 105:
-    batch_foreach<tensor_foreach<kernel_scale_transverse_mass<CpuTypes>, kernel_scale_transverse_mass<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_scale_transverse_energy<CpuTypes>, kernel_scale_transverse_energy<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 106:
-    batch_foreach<tensor_foreach<kernel_scale_half_transverse_mass<CpuTypes>, kernel_scale_half_transverse_mass<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_scale_transverse_mass<CpuTypes>, kernel_scale_transverse_mass<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 107:
-    batch_foreach<tensor_foreach<kernel_scale_partonic_energy<CpuTypes>, kernel_scale_partonic_energy<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_scale_half_transverse_mass<CpuTypes>, kernel_scale_half_transverse_mass<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 108:
-    batch_foreach<tensor_foreach<kernel_chili_forward<CpuTypes>, kernel_chili_forward<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_scale_partonic_energy<CpuTypes>, kernel_scale_partonic_energy<SimdTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 109:
-    batch_foreach<tensor_foreach<kernel_chili_inverse<CpuTypes>, kernel_chili_inverse<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_chili_forward<CpuTypes>, kernel_chili_forward<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
     break;
 case 110:
-    op_matrix_element(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_chili_inverse<CpuTypes>, kernel_chili_inverse<SimdTypes>, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device);
     break;
 case 111:
-    batch_foreach<tensor_foreach<kernel_collect_channel_weights<CpuTypes>, kernel_collect_channel_weights<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
+    op_matrix_element(instr, locals, device);
     break;
 case 112:
-    batch_foreach<tensor_foreach<kernel_interpolate_pdf<CpuTypes>, kernel_interpolate_pdf<CpuTypes>, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_collect_channel_weights<CpuTypes>, kernel_collect_channel_weights<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
     break;
 case 113:
-    batch_foreach<tensor_foreach<kernel_interpolate_alpha_s<CpuTypes>, kernel_interpolate_alpha_s<CpuTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_interpolate_pdf<CpuTypes>, kernel_interpolate_pdf<CpuTypes>, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device);
     break;
 case 114:
-    op_matmul(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_interpolate_alpha_s<CpuTypes>, kernel_interpolate_alpha_s<CpuTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
     break;
 case 115:
-    batch_foreach<tensor_foreach_dynamic<kernel_relu<CpuTypes>, kernel_relu<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    op_matmul(instr, locals, device);
     break;
 case 116:
-    batch_foreach<tensor_foreach_dynamic<kernel_leaky_relu<CpuTypes>, kernel_leaky_relu<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_relu<CpuTypes>, kernel_relu<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 117:
-    batch_foreach<tensor_foreach_dynamic<kernel_elu<CpuTypes>, kernel_elu<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_leaky_relu<CpuTypes>, kernel_leaky_relu<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 118:
-    batch_foreach<tensor_foreach_dynamic<kernel_gelu<CpuTypes>, kernel_gelu<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_elu<CpuTypes>, kernel_elu<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 119:
-    batch_foreach<tensor_foreach_dynamic<kernel_sigmoid<CpuTypes>, kernel_sigmoid<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_gelu<CpuTypes>, kernel_gelu<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 120:
-    batch_foreach<tensor_foreach_dynamic<kernel_softplus<CpuTypes>, kernel_softplus<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_sigmoid<CpuTypes>, kernel_sigmoid<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 121:
-    op_rqs_reshape(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_softplus<CpuTypes>, kernel_softplus<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 122:
-    batch_foreach<tensor_foreach<kernel_rqs_find_bin<CpuTypes>, kernel_rqs_find_bin<SimdTypes>, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device);
+    op_rqs_reshape(instr, locals, device);
     break;
 case 123:
-    batch_foreach<tensor_foreach<kernel_rqs_forward<CpuTypes>, kernel_rqs_forward<SimdTypes>, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_rqs_find_bin<CpuTypes>, kernel_rqs_find_bin<SimdTypes>, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device);
     break;
 case 124:
-    batch_foreach<tensor_foreach<kernel_rqs_inverse<CpuTypes>, kernel_rqs_inverse<SimdTypes>, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_rqs_forward<CpuTypes>, kernel_rqs_forward<SimdTypes>, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device);
     break;
 case 125:
-    batch_foreach<tensor_foreach_dynamic<kernel_softmax<CpuTypes>, kernel_softmax<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_rqs_inverse<CpuTypes>, kernel_rqs_inverse<SimdTypes>, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device);
     break;
 case 126:
-    batch_foreach<tensor_foreach<kernel_softmax_prior<CpuTypes>, kernel_softmax_prior<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_softmax<CpuTypes>, kernel_softmax<SimdTypes>, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 127:
-    batch_foreach<tensor_foreach<kernel_sample_discrete<CpuTypes>, kernel_sample_discrete<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_softmax_prior<CpuTypes>, kernel_softmax_prior<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 128:
-    batch_foreach<tensor_foreach<kernel_sample_discrete_inverse<CpuTypes>, kernel_sample_discrete_inverse<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_sample_discrete<CpuTypes>, kernel_sample_discrete<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
     break;
 case 129:
-    batch_foreach<tensor_foreach<kernel_sample_discrete_probs<CpuTypes>, kernel_sample_discrete_probs<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_sample_discrete_inverse<CpuTypes>, kernel_sample_discrete_inverse<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
     break;
 case 130:
-    batch_foreach<tensor_foreach<kernel_sample_discrete_probs_inverse<CpuTypes>, kernel_sample_discrete_probs_inverse<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_sample_discrete_probs<CpuTypes>, kernel_sample_discrete_probs<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
     break;
 case 131:
-    op_discrete_histogram(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_sample_discrete_probs_inverse<CpuTypes>, kernel_sample_discrete_probs_inverse<SimdTypes>, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device);
     break;
 case 132:
-    batch_foreach<tensor_foreach<kernel_permute_momenta<CpuTypes>, kernel_permute_momenta<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
+    op_discrete_histogram(instr, locals, device);
     break;
 case 133:
-    batch_foreach<tensor_foreach<kernel_gather<CpuTypes>, kernel_gather<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_permute_momenta<CpuTypes>, kernel_permute_momenta<SimdTypes>, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device);
     break;
 case 134:
-    batch_foreach<tensor_foreach<kernel_gather_int<CpuTypes>, kernel_gather_int<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_gather<CpuTypes>, kernel_gather<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 135:
-    batch_foreach<tensor_foreach<kernel_gather_vector<CpuTypes>, kernel_gather_vector<CpuTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_gather_int<CpuTypes>, kernel_gather_int<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 136:
-    batch_foreach<tensor_foreach<kernel_select_int<CpuTypes>, kernel_select_int<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_gather_vector<CpuTypes>, kernel_gather_vector<CpuTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 137:
-    batch_foreach<tensor_foreach<kernel_select<CpuTypes>, kernel_select<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_select_int<CpuTypes>, kernel_select_int<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 138:
-    batch_foreach<tensor_foreach<kernel_select_vector<CpuTypes>, kernel_select_vector<CpuTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_select<CpuTypes>, kernel_select<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 139:
-    batch_foreach<tensor_foreach<kernel_argsort<CpuTypes>, kernel_argsort<CpuTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_select_vector<CpuTypes>, kernel_select_vector<CpuTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 140:
-    op_quantile(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_argsort<CpuTypes>, kernel_argsort<CpuTypes>, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device);
     break;
 case 141:
-    batch_foreach<tensor_foreach<kernel_one_hot<CpuTypes>, kernel_one_hot<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    op_quantile(instr, locals, device);
     break;
 case 142:
-    batch_foreach<tensor_foreach<kernel_madnis_abs_weight<CpuTypes>, kernel_madnis_abs_weight<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_one_hot<CpuTypes>, kernel_one_hot<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 143:
-    batch_foreach<tensor_foreach<kernel_madnis_softclip<CpuTypes>, kernel_madnis_softclip<SimdTypes>, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_madnis_abs_weight<CpuTypes>, kernel_madnis_abs_weight<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 144:
-    batch_foreach<tensor_foreach<kernel_madnis_variance<CpuTypes>, kernel_madnis_variance<SimdTypes>, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_madnis_softclip<CpuTypes>, kernel_madnis_softclip<SimdTypes>, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device);
     break;
 case 145:
-    batch_foreach<tensor_foreach<kernel_madnis_single_channel_variance<CpuTypes>, kernel_madnis_single_channel_variance<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_madnis_variance<CpuTypes>, kernel_madnis_variance<SimdTypes>, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device);
     break;
 case 146:
-    batch_foreach<tensor_foreach<kernel_madnis_multi_channel_variance<CpuTypes>, kernel_madnis_multi_channel_variance<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_madnis_single_channel_variance<CpuTypes>, kernel_madnis_single_channel_variance<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 147:
-    op_nonzero(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_madnis_multi_channel_variance<CpuTypes>, kernel_madnis_multi_channel_variance<SimdTypes>, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device);
     break;
 case 148:
-    op_batch_gather(instr, locals, device);
+    op_nonzero(instr, locals, device);
     break;
 case 149:
-    op_batch_scatter(instr, locals, device);
+    op_batch_gather(instr, locals, device);
     break;
 case 150:
-    op_random(instr, locals, device);
+    op_batch_scatter(instr, locals, device);
     break;
 case 151:
-    op_random_int(instr, locals, device);
+    op_random(instr, locals, device);
     break;
 case 152:
-    op_unweight(instr, locals, device);
+    op_random_int(instr, locals, device);
     break;
 case 153:
-    batch_foreach<tensor_foreach<kernel_vegas_forward<CpuTypes>, kernel_vegas_forward<SimdTypes>, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device);
+    op_unweight(instr, locals, device);
     break;
 case 154:
-    batch_foreach<tensor_foreach<kernel_vegas_inverse<CpuTypes>, kernel_vegas_inverse<SimdTypes>, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_vegas_forward<CpuTypes>, kernel_vegas_forward<SimdTypes>, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device);
     break;
 case 155:
-    op_vegas_histogram(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_vegas_inverse<CpuTypes>, kernel_vegas_inverse<SimdTypes>, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device);
     break;
 case 156:
+    op_vegas_histogram(instr, locals, device);
+    break;
+case 157:
     op_histogram(instr, locals, device);
     break;
diff --git a/madspace/src/driver/adam_optimizer.cpp b/madspace/src/driver/adam_optimizer.cpp
index e7959cf695..c3686e8f70 100644
--- a/madspace/src/driver/adam_optimizer.cpp
+++ b/madspace/src/driver/adam_optimizer.cpp
@@ -4,6 +4,27 @@
 
 using namespace madspace;
 
+GradientClipper::GradientClipper(double threshold) :
+    FunctionGenerator(
+        "Unweighter",
+        {{"gradients_in", batch_float}, {"threshold", single_float}},
+        {{"gradients_out", batch_float}, {"gradient_norm", single_float}}
+    ),
+    _threshold(threshold) {}
+
+NamedVector<Value> GradientClipper::build_function_impl(
+    FunctionBuilder& fb, const NamedVector<Value>& args
+) const {
+    Value grads_in = args.at(0);
+    Value grad_norm = fb.sqrt(fb.batch_reduce_sum(fb.square(grads_in)));
+    Value factor = fb.min(fb.div(_threshold, grad_norm), 1.0);
+    Value grads_out = fb.mul(grads_in, factor);
+    return {
+        {"gradients_out", grads_out},
+        {"gradient_norm", grad_norm},
+    };
+}
+
 AdamOptimizer::AdamOptimizer(
     const Function& function,
     ContextPtr context,
@@ -12,7 +33,8 @@ AdamOptimizer::AdamOptimizer(
     std::size_t step_count,
     double beta1,
     double beta2,
-    double eps
+    double eps,
+    double grad_clip_threshold
 ) :
     _context(context),
     _learning_rate(learning_rate),
@@ -22,6 +44,13 @@ AdamOptimizer::AdamOptimizer(
     _beta1(beta1),
     _beta2(beta2),
     _eps(eps),
+    _grad_clipper(
+        grad_clip_threshold > 0.
+            ? build_runtime(
+                  GradientClipper(grad_clip_threshold).function(), context, false
+              )
+            : nullptr
+    ),
     _one(1.0, context->device()) {
     DevicePtr device = context->device();
     for (auto& [name, value] : function.globals()) {
@@ -61,6 +90,11 @@ TensorVec AdamOptimizer::step(const TensorVec& inputs) {
     output_grads.at(0) = _one;
     auto [input_grads, global_grads] =
         _runtime->run_backward(output_grads, stored_locals, eval_grad, true);
+
+    if (_grad_clipper) {
+        global_grads = {_grad_clipper->run(global_grads).at(0)};
+    }
+
     device->adam_step(
         global_grads.at(0),
         _parameter,
diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp
index a5e4760a5e..dcf99ae523 100644
--- a/madspace/src/driver/event_generator.cpp
+++ b/madspace/src/driver/event_generator.cpp
@@ -1,7 +1,6 @@
 #include "madspace/driver/event_generator.hpp"
 
 #include <cmath>
-#include <filesystem>
 #include <format>
 #include <ranges>
 
@@ -15,7 +14,7 @@ const GeneratorConfig EventGenerator::default_config = {};
 EventGenerator::EventGenerator(
     const std::vector<ContextPtr>& contexts,
     const std::vector<std::shared_ptr<ChannelEventGenerator>>& channels,
-    const std::string& status_file,
+    std::shared_ptr<StatusFile> status_file,
     const GeneratorConfig& config
 ) :
     _config(config),
@@ -723,20 +722,13 @@ void EventGenerator::fill_lhe_event(
 }
 
 void EventGenerator::init_status(const std::string& status) {
-    _last_status_time = std::chrono::steady_clock::now();
     write_status(status, true);
 }
 
 void EventGenerator::write_status(const std::string& status, bool force_write) {
-    auto now = std::chrono::steady_clock::now();
-    using namespace std::chrono_literals;
-    if (now - _last_status_time < 10s && !force_write) {
+    if (!_status_file) {
         return;
     }
-    _last_status_time = now;
-
-    std::string status_tmp_file = std::format("{}.tmp", _status_file);
-    std::ofstream f(status_tmp_file);
     nlohmann::json j{
         {"status", status},
         {"process", _status},
@@ -744,10 +736,7 @@ void EventGenerator::write_status(const std::string& status, bool force_write) {
         {"run_times", _timing_data},
         {"histograms", histograms()},
     };
-    f << j.dump();
-    // rename atomically deletes the old file and replaces it with the new one
-    // such that the status file exists at all times
-    std::filesystem::rename(status_tmp_file, _status_file);
+    _status_file->write(j, force_write);
 }
 
 void EventGenerator::print_survey_init() {
diff --git a/madspace/src/driver/madnis_training.cpp b/madspace/src/driver/madnis_training.cpp
index c94d84fb94..470174d8f9 100644
--- a/madspace/src/driver/madnis_training.cpp
+++ b/madspace/src/driver/madnis_training.cpp
@@ -53,11 +53,13 @@ void MadnisTraining::train_step(std::size_t batch_index) {
     bool try_buffered =
         _config.buffer_capacity > 0 && batch_index % (_config.buffered_steps + 1) != 0;
     TensorVec training_batch;
+    bool used_buffered = false;
     while (true) {
         start_generator_jobs(channel_sizes);
         if (try_buffered) {
             if (check_buffered_training_batch(channel_sizes)) {
                 training_batch = build_buffered_training_batch(channel_sizes);
+                used_buffered = true;
                 break;
             }
             try_buffered = false;
@@ -67,8 +69,18 @@ void MadnisTraining::train_step(std::size_t batch_index) {
         }
         process_job_results(gen_thread_pool.wait_multiple());
     }
+    double learning_rate = _optimizer->learning_rate();
     TensorVec results = _optimizer->step(training_batch);
-    update_history(results, channel_sizes);
+    update_history(results, channel_sizes, learning_rate, used_buffered);
+    if ((batch_index + 1) % _config.log_interval == 0) {
+        _status_batches.push_back(batch_index + 1);
+        _status_losses.push_back(average_loss());
+        _status_channel_counts.push_back(active_channel_count());
+        _status_learning_rates.push_back(average_learning_rate());
+        _status_buffered_fractions.push_back(buffered_fraction());
+        _status_generated_events.push_back(_generated_event_count);
+        _status_buffer_sizes.push_back(buffer_event_count());
+    }
     if (_channels.size() > 0 && _cwnet &&
         (batch_index + 1) % _config.channel_dropping_interval == 0) {
         std::vector<std::size_t> job_ids;
@@ -111,6 +123,36 @@ double MadnisTraining::average_loss() const {
     return loss_sum / loss_count;
 }
 
+double MadnisTraining::average_learning_rate() const {
+    if (_lr_history.size() == 0) {
+        return 0.;
+    }
+    double lr_sum = 0;
+    for (double lr : _lr_history) {
+        lr_sum += lr;
+    }
+    return lr_sum / _lr_history.size();
+}
+
+double MadnisTraining::buffered_fraction() const {
+    if (_buffered_history.size() == 0) {
+        return 0.;
+    }
+    std::size_t buffered_count = 0;
+    for (bool buffered : _buffered_history) {
+        buffered_count += buffered;
+    }
+    return static_cast<double>(buffered_count) / _buffered_history.size();
+}
+
+std::size_t MadnisTraining::buffer_event_count() const {
+    std::size_t count = 0;
+    for (auto& channel : _channels) {
+        count += channel.buffer.size;
+    }
+    return count;
+}
+
 void MadnisTraining::build_runtimes_and_optimizer() {
     std::vector<std::shared_ptr<FunctionGenerator>> functions;
     functions.reserve(_channels.size());
@@ -132,7 +174,8 @@ void MadnisTraining::build_runtimes_and_optimizer() {
             _config.batches,
             _config.adam_beta1,
             _config.adam_beta2,
-            _config.adam_eps
+            _config.adam_eps,
+            _config.grad_clip_threshold
         );
     }
     _generator_params =
@@ -511,6 +554,7 @@ void MadnisTraining::process_job_results(const std::vector<std::size_t>& job_ids
         if (job.samples.channel_sizes.size() == 0) {
             auto& channel = _channels.at(job.samples.channel_index);
             channel.sample_count += job.samples.size;
+            _generated_event_count += job.samples.size;
             channel.sample_batches.push_back(std::move(job.samples));
             if (job.unweighted_samples.size > 0) {
                 buffer_store(channel, job.unweighted_samples);
@@ -525,6 +569,7 @@ void MadnisTraining::process_job_results(const std::vector<std::size_t>& job_ids
                     continue;
                 }
                 channel.sample_count += chan_size;
+                _generated_event_count += chan_size;
                 channel.sample_batches.emplace_back();
                 auto& batch = channel.sample_batches.back();
                 batch.tensors.reserve(job.samples.tensors.size());
@@ -589,16 +634,26 @@ void MadnisTraining::buffer_store(ChannelData& channel, SampleBatch& samples) {
 }
 
 void MadnisTraining::update_history(
-    const TensorVec& results, const std::vector<std::size_t>& counts
+    const TensorVec& results,
+    const std::vector<std::size_t>& counts,
+    double learning_rate,
+    bool buffered
 ) {
     Tensor loss_cpu = results.at(0).cpu();
     Tensor abs_means_cpu = results.at(1).cpu();
     Tensor variances_cpu = results.at(2).cpu();
     double loss = loss_cpu.view<double, 1>()[0];
+    if (loss > 1e6) {
+        throw std::runtime_error("MadNIS training diverged. Please restart");
+    }
     if (_loss_history.size() < _config.log_interval) {
         _loss_history.push_back(loss);
+        _lr_history.push_back(learning_rate);
+        _buffered_history.push_back(buffered);
     } else {
         _loss_history.at(_loss_history_index) = loss;
+        _lr_history.at(_loss_history_index) = learning_rate;
+        _buffered_history.at(_loss_history_index) = buffered;
     }
     if (++_loss_history_index == _config.log_interval) {
         _loss_history_index = 0;
@@ -719,9 +774,10 @@ MultiMadnisTraining::MultiMadnisTraining(
     ContextPtr optimizer_context,
     const MadnisTraining::Config& config,
     const nested_vector2<std::shared_ptr<Integrand>>& integrands,
-    const std::vector<std::optional<ChannelWeightNetwork>>& cwnets
+    const std::vector<std::optional<ChannelWeightNetwork>>& cwnets,
+    std::shared_ptr<StatusFile> status_file
 ) :
-    _config(config) {
+    _config(config), _status_file(status_file) {
     _subprocesses.reserve(integrands.size());
     for (auto [integ, cwnet] : zip(integrands, cwnets)) {
         _subprocesses.emplace_back(
@@ -741,6 +797,9 @@ void MultiMadnisTraining::train() {
             double loss = subproc.average_loss();
             std::size_t chan_count = subproc.active_channel_count();
             print_progress_update(subproc_index, batch_index, loss, chan_count);
+            bool done = subproc_index == _subprocesses.size() - 1 &&
+                batch_index + 1 == _config.batches;
+            write_status(subproc_index, batch_index, done);
         }
         ++subproc_index;
     }
@@ -887,3 +946,43 @@ void MultiMadnisTraining::print_progress_update(
         }
     }
 }
+
+void MultiMadnisTraining::write_status(
+    std::size_t subproc_index, std::size_t batch_index, bool done
+) {
+    if (!_status_file) {
+        return;
+    }
+    if (!done && (batch_index + 1) % _config.log_interval != 0) {
+        return;
+    }
+    using namespace std::chrono_literals;
+    auto now = std::chrono::steady_clock::now();
+
+    nlohmann::json trainings = nlohmann::json::array();
+    for (std::size_t i = 0; auto& subproc : _subprocesses) {
+        trainings.push_back(
+            {{"subprocess", i},
+             {"batch", subproc.status_batches()},
+             {"batch_count", _config.batches},
+             {"losses", subproc.status_losses()},
+             {"channel_counts", subproc.status_channel_counts()},
+             {"learning_rates", subproc.status_learning_rates()},
+             {"buffered_fractions", subproc.status_buffered_fractions()},
+             {"generated_events", subproc.status_generated_events()},
+             {"buffer_sizes", subproc.status_buffer_sizes()}}
+        );
+        ++i;
+    }
+
+    nlohmann::json j{
+        {"status", done ? "done" : "training"},
+        {"madnis_trainings", trainings},
+        {"run_times",
+         {{"training",
+           {{"wall_time_sec", (now - _start_time) / 1.0s},
+            {"cpu_time_sec", (cpu_time_microsec() - _start_cpu_microsec) / 1e6}}}}},
+    };
+    bool force_write = done || (subproc_index == 0 && batch_index == 0);
+    _status_file->write(j, force_write);
+}
diff --git a/madspace/src/driver/status_file.cpp b/madspace/src/driver/status_file.cpp
new file mode 100644
index 0000000000..8ece6ed931
--- /dev/null
+++ b/madspace/src/driver/status_file.cpp
@@ -0,0 +1,44 @@
+#include "madspace/driver/status_file.hpp"
+
+#include <filesystem>
+#include <format>
+#include <fstream>
+
+using namespace madspace;
+
+StatusFile::StatusFile(const std::string& file_name, double min_interval_sec) :
+    _file_name(file_name), _min_interval(min_interval_sec) {}
+
+void StatusFile::write(const nlohmann::json& content, bool force_write) {
+    for (auto& [key, value] : content.items()) {
+        if (key == "run_times" && value.is_object()) {
+            auto& run_times = _content["run_times"];
+            if (!run_times.is_object()) {
+                run_times = nlohmann::json::object();
+            }
+            run_times.update(value);
+        } else {
+            _content[key] = value;
+        }
+    }
+
+    auto now = std::chrono::steady_clock::now();
+    if (_has_written && !force_write && now - _last_write_time < _min_interval) {
+        return;
+    }
+    _has_written = true;
+    _last_write_time = now;
+
+    std::string status_tmp_file = std::format("{}.tmp", _file_name);
+    std::ofstream f(status_tmp_file);
+    f << _content.dump();
+    // rename atomically deletes the old file and replaces it with the new one
+    // such that the status file exists at all times
+    std::filesystem::rename(status_tmp_file, _file_name);
+}
+
+StatusFile::~StatusFile() {
+    if (_has_written && _content.value("status", "") != "done") {
+        write({{"status", "done"}}, true);
+    }
+}
diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu
index 87ff0b0bec..af9dde25d0 100644
--- a/madspace/src/gpu/runtime.cu
+++ b/madspace/src/gpu/runtime.cu
@@ -546,11 +546,12 @@ __global__ void kernel_div_batch_size(
     }
 }
 
-void batch_reduce_mean_impl(
+void batch_reduce_sum_mean_impl(
     const GpuRuntime::Instruction& instruction,
     TensorVec& locals,
     const AsyncGpuDevice& device,
-    bool keepdim
+    bool keepdim,
+    bool is_mean
 ) {
     auto input = locals[instruction.input_indices[0]].contiguous(device);
     auto& output = locals[instruction.output_indices[0]];
@@ -580,14 +581,16 @@ void batch_reduce_mean_impl(
     );
     temp.reset(device);
     input.reset(device);
-    launch_kernel(
-        kernel_div_batch_size,
-        1,
-        device.stream(),
-        batch_size,
-        out.view<double, 1>(),
-        out.view<double, 1>()
-    );
+    if (is_mean) {
+        launch_kernel(
+            kernel_div_batch_size,
+            1,
+            device.stream(),
+            batch_size,
+            out.view<double, 1>(),
+            out.view<double, 1>()
+        );
+    }
     if (keepdim) {
         output = out.expand({batch_size});
     } else {
@@ -595,12 +598,13 @@ void batch_reduce_mean_impl(
     }
 }
 
-void batch_reduce_mean_backward_impl(
+void batch_reduce_sum_mean_backward_impl(
     const GpuRuntime::Instruction& instruction,
     TensorVec& locals,
     TensorVec& local_grads,
     const AsyncGpuDevice& device,
-    bool keepdim
+    bool keepdim,
+    bool is_mean
 ) {
     auto& input = locals[instruction.input_indices[0]];
     auto& input_grad = local_grads[instruction.input_indices[0]];
@@ -642,37 +646,62 @@ void batch_reduce_mean_backward_impl(
             batch_size,
             device.stream()
         );
-        launch_kernel(
-            kernel_div_batch_size,
-            1,
-            device.stream(),
-            output_grad.size(0),
-            grad.view<double, 1>(),
-            grad.view<double, 1>()
-        );
         temp.reset(device);
         output_grad.reset(device);
+        if (is_mean) {
+            launch_kernel(
+                kernel_div_batch_size,
+                1,
+                device.stream(),
+                batch_size,
+                grad.view<double, 1>(),
+                grad.view<double, 1>()
+            );
+        }
     } else {
         auto& output_grad = local_grads[instruction.output_indices[0]];
-        launch_kernel(
-            kernel_div_batch_size,
-            1,
-            device.stream(),
-            output_grad.size(0),
-            output_grad.view<double, 1>(),
-            grad.view<double, 1>()
-        );
+        if (is_mean) {
+            launch_kernel(
+                kernel_div_batch_size,
+                1,
+                device.stream(),
+                input.size(0),
+                output_grad.view<double, 1>(),
+                grad.view<double, 1>()
+            );
+        } else {
+            grad.copy_from(output_grad, device);
+        }
     }
     input_grad.add(grad, device);
     grad.reset(device);
 }
 
+void op_batch_reduce_sum(
+    const GpuRuntime::Instruction& instruction,
+    TensorVec& locals,
+    const AsyncGpuDevice& device
+) {
+    batch_reduce_sum_mean_impl(instruction, locals, device, false, false);
+}
+
+void backward_op_batch_reduce_sum(
+    const GpuRuntime::Instruction& instruction,
+    TensorVec& locals,
+    TensorVec& local_grads,
+    const AsyncGpuDevice& device
+) {
+    batch_reduce_sum_mean_backward_impl(
+        instruction, locals, local_grads, device, false, false
+    );
+}
+
 void op_batch_reduce_mean(
     const GpuRuntime::Instruction& instruction,
     TensorVec& locals,
     const AsyncGpuDevice& device
 ) {
-    batch_reduce_mean_impl(instruction, locals, device, false);
+    batch_reduce_sum_mean_impl(instruction, locals, device, false, true);
 }
 
 void backward_op_batch_reduce_mean(
@@ -681,7 +710,9 @@ void backward_op_batch_reduce_mean(
     TensorVec& local_grads,
     const AsyncGpuDevice& device
 ) {
-    batch_reduce_mean_backward_impl(instruction, locals, local_grads, device, false);
+    batch_reduce_sum_mean_backward_impl(
+        instruction, locals, local_grads, device, false, true
+    );
 }
 
 void op_batch_reduce_mean_keepdim(
@@ -689,7 +720,7 @@ void op_batch_reduce_mean_keepdim(
     TensorVec& locals,
     const AsyncGpuDevice& device
 ) {
-    batch_reduce_mean_impl(instruction, locals, device, true);
+    batch_reduce_sum_mean_impl(instruction, locals, device, true, true);
 }
 
 void backward_op_batch_reduce_mean_keepdim(
@@ -698,7 +729,9 @@ void backward_op_batch_reduce_mean_keepdim(
     TensorVec& local_grads,
     const AsyncGpuDevice& device
 ) {
-    batch_reduce_mean_backward_impl(instruction, locals, local_grads, device, true);
+    batch_reduce_sum_mean_backward_impl(
+        instruction, locals, local_grads, device, true, true
+    );
 }
 
 void op_offset_indices(
diff --git a/madspace/src/gpu/runtime_backward_mixin.inc b/madspace/src/gpu/runtime_backward_mixin.inc
index 452946c959..3f718287d7 100644
--- a/madspace/src/gpu/runtime_backward_mixin.inc
+++ b/madspace/src/gpu/runtime_backward_mixin.inc
@@ -38,80 +38,83 @@ case 18:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_div<GpuTypes>, 3, 2>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
 case 21:
-    backward_op_batch_reduce_mean(instr, locals, local_grads, device);
+    backward_op_batch_reduce_sum(instr, locals, local_grads, device);
     break;
 case 22:
-    backward_op_batch_reduce_mean_keepdim(instr, locals, local_grads, device);
+    backward_op_batch_reduce_mean(instr, locals, local_grads, device);
     break;
 case 23:
-    backward_batch_foreach<tensor_foreach<backward_kernel_reduce_product<GpuTypes>, 2, 1, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
+    backward_op_batch_reduce_mean_keepdim(instr, locals, local_grads, device);
     break;
 case 24:
-    backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_sqrt<GpuTypes>, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
+    backward_batch_foreach<tensor_foreach<backward_kernel_reduce_product<GpuTypes>, 2, 1, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
 case 25:
+    backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_sqrt<GpuTypes>, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
+    break;
+case 26:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_square<GpuTypes>, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 114:
+case 115:
     backward_op_matmul(instr, locals, local_grads, device);
     break;
-case 115:
+case 116:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_relu<GpuTypes>, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 116:
+case 117:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_leaky_relu<GpuTypes>, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 117:
+case 118:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_elu<GpuTypes>, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 118:
+case 119:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_gelu<GpuTypes>, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 119:
+case 120:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_sigmoid<GpuTypes>, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device);
     break;
-case 120:
+case 121:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_softplus<GpuTypes>, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 121:
+case 122:
     backward_op_rqs_reshape(instr, locals, local_grads, device);
     break;
-case 122:
+case 123:
     backward_batch_foreach<tensor_foreach<backward_kernel_rqs_find_bin<GpuTypes>, 5, 4, 2>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device);
     break;
-case 123:
+case 124:
     backward_batch_foreach<tensor_foreach<backward_kernel_rqs_forward<GpuTypes>, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
-case 124:
+case 125:
     backward_batch_foreach<tensor_foreach<backward_kernel_rqs_inverse<GpuTypes>, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
-case 125:
+case 126:
     backward_batch_foreach<tensor_foreach_dynamic<backward_kernel_softmax<GpuTypes>, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device);
     break;
-case 126:
+case 127:
     backward_batch_foreach<tensor_foreach<backward_kernel_softmax_prior<GpuTypes>, 2, 2, 1>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device);
     break;
-case 130:
+case 131:
     backward_batch_foreach<tensor_foreach<backward_kernel_sample_discrete_probs_inverse<GpuTypes>, 4, 2, 1>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
-case 133:
+case 134:
     backward_batch_foreach<tensor_foreach<backward_kernel_gather<GpuTypes>, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device);
     break;
-case 137:
+case 138:
     backward_batch_foreach<tensor_foreach<backward_kernel_select<GpuTypes>, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device);
     break;
-case 142:
+case 143:
     backward_batch_foreach<tensor_foreach<backward_kernel_madnis_abs_weight<GpuTypes>, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
-case 143:
+case 144:
     backward_batch_foreach<tensor_foreach<backward_kernel_madnis_softclip<GpuTypes>, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device);
     break;
-case 144:
+case 145:
     backward_batch_foreach<tensor_foreach<backward_kernel_madnis_variance<GpuTypes>, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device);
     break;
-case 145:
+case 146:
     backward_batch_foreach<tensor_foreach<backward_kernel_madnis_single_channel_variance<GpuTypes>, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device);
     break;
-case 146:
+case 147:
     backward_batch_foreach<tensor_foreach<backward_kernel_madnis_multi_channel_variance<GpuTypes>, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device);
     break;
diff --git a/madspace/src/gpu/runtime_mixin.inc b/madspace/src/gpu/runtime_mixin.inc
index 316c746a3d..6ec6132815 100644
--- a/madspace/src/gpu/runtime_mixin.inc
+++ b/madspace/src/gpu/runtime_mixin.inc
@@ -65,410 +65,413 @@ case 20:
     batch_foreach<tensor_foreach_dynamic<kernel_reduce_sum_vector<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 21:
-    op_batch_reduce_mean(instr, locals, device);
+    op_batch_reduce_sum(instr, locals, device);
     break;
 case 22:
-    op_batch_reduce_mean_keepdim(instr, locals, device);
+    op_batch_reduce_mean(instr, locals, device);
     break;
 case 23:
-    batch_foreach<tensor_foreach<kernel_reduce_product<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
+    op_batch_reduce_mean_keepdim(instr, locals, device);
     break;
 case 24:
-    batch_foreach<tensor_foreach_dynamic<kernel_sqrt<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_reduce_product<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 25:
-    batch_foreach<tensor_foreach_dynamic<kernel_square<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_sqrt<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 26:
-    batch_foreach<tensor_foreach_dynamic<kernel_min<GpuTypes>, 2, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_square<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 27:
-    batch_foreach<tensor_foreach_dynamic<kernel_max<GpuTypes>, 2, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_min<GpuTypes>, 2, 1>, 2, 1>(instr, locals, device);
     break;
 case 28:
-    batch_foreach<tensor_foreach<kernel_obs_sqrt_s<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_max<GpuTypes>, 2, 1>, 2, 1>(instr, locals, device);
     break;
 case 29:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_e<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_obs_sqrt_s<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 30:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_px<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_e<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 31:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_py<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_px<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 32:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_pz<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_py<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 33:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_mass<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_pz<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 34:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_pt<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_mass<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 35:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_p_mag<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_pt<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 36:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_phi<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_p_mag<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 37:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_theta<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_phi<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 38:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_y<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_theta<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 39:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_y_abs<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_y<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 40:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_eta<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_y_abs<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 41:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_eta_abs<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_eta<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 42:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_eta<GpuTypes>, 2, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_eta_abs<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 43:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_phi<GpuTypes>, 2, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_eta<GpuTypes>, 2, 1>, 2, 1>(instr, locals, device);
     break;
 case 44:
-    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_r<GpuTypes>, 2, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_phi<GpuTypes>, 2, 1>, 2, 1>(instr, locals, device);
     break;
 case 45:
-    batch_foreach<tensor_foreach<kernel_boost_beam<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_obs_delta_r<GpuTypes>, 2, 1>, 2, 1>(instr, locals, device);
     break;
 case 46:
-    batch_foreach<tensor_foreach<kernel_boost_beam_inverse<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_boost_beam<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
     break;
 case 47:
-    batch_foreach<tensor_foreach<kernel_com_p_in<GpuTypes>, 1, 2, 1>, 1, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_boost_beam_inverse<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
     break;
 case 48:
-    batch_foreach<tensor_foreach<kernel_r_to_x1x2<GpuTypes>, 3, 3, 1>, 3, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_com_p_in<GpuTypes>, 1, 2, 1>, 1, 2>(instr, locals, device);
     break;
 case 49:
-    batch_foreach<tensor_foreach<kernel_x1x2_to_r<GpuTypes>, 3, 2, 1>, 3, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_r_to_x1x2<GpuTypes>, 3, 3, 1>, 3, 3>(instr, locals, device);
     break;
 case 50:
-    batch_foreach<tensor_foreach<kernel_diff_cross_section<GpuTypes>, 6, 1, 1>, 6, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_x1x2_to_r<GpuTypes>, 3, 2, 1>, 3, 2>(instr, locals, device);
     break;
 case 51:
-    batch_foreach<tensor_foreach<kernel_two_body_decay_com<GpuTypes>, 5, 3, 1>, 5, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_diff_cross_section<GpuTypes>, 6, 1, 1>, 6, 1>(instr, locals, device);
     break;
 case 52:
-    batch_foreach<tensor_foreach<kernel_two_body_decay_com_inverse<GpuTypes>, 2, 6, 1>, 2, 6>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_body_decay_com<GpuTypes>, 5, 3, 1>, 5, 3>(instr, locals, device);
     break;
 case 53:
-    batch_foreach<tensor_foreach<kernel_two_body_decay<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_body_decay_com_inverse<GpuTypes>, 2, 6, 1>, 2, 6>(instr, locals, device);
     break;
 case 54:
-    batch_foreach<tensor_foreach<kernel_two_body_decay_inverse<GpuTypes>, 2, 7, 1>, 2, 7>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_body_decay<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
     break;
 case 55:
-    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_com<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_body_decay_inverse<GpuTypes>, 2, 7, 1>, 2, 7>(instr, locals, device);
     break;
 case 56:
-    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_com_inverse<GpuTypes>, 4, 4, 1>, 4, 4>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_com<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
     break;
 case 57:
-    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_com_inverse<GpuTypes>, 4, 4, 1>, 4, 4>(instr, locals, device);
     break;
 case 58:
-    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_inverse<GpuTypes>, 4, 4, 1>, 4, 4>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
     break;
 case 59:
-    batch_foreach<tensor_foreach<kernel_two_to_three_particle_scattering<GpuTypes>, 8, 3, 1>, 8, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_two_particle_scattering_inverse<GpuTypes>, 4, 4, 1>, 4, 4>(instr, locals, device);
     break;
 case 60:
-    batch_foreach<tensor_foreach<kernel_two_to_three_particle_scattering_inverse<GpuTypes>, 7, 4, 1>, 7, 4>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_three_particle_scattering<GpuTypes>, 8, 3, 1>, 8, 3>(instr, locals, device);
     break;
 case 61:
-    batch_foreach<tensor_foreach<kernel_double_t_scattering<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_two_to_three_particle_scattering_inverse<GpuTypes>, 7, 4, 1>, 7, 4>(instr, locals, device);
     break;
 case 62:
-    batch_foreach<tensor_foreach<kernel_double_t_scattering_inverse<GpuTypes>, 4, 2, 1>, 4, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_double_t_scattering<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
     break;
 case 63:
-    batch_foreach<tensor_foreach<kernel_three_body_decay_com<GpuTypes>, 9, 4, 1>, 9, 4>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_double_t_scattering_inverse<GpuTypes>, 4, 2, 1>, 4, 2>(instr, locals, device);
     break;
 case 64:
-    batch_foreach<tensor_foreach<kernel_three_body_decay_com_inverse<GpuTypes>, 3, 10, 1>, 3, 10>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_three_body_decay_com<GpuTypes>, 9, 4, 1>, 9, 4>(instr, locals, device);
     break;
 case 65:
-    batch_foreach<tensor_foreach<kernel_three_body_decay<GpuTypes>, 10, 4, 1>, 10, 4>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_three_body_decay_com_inverse<GpuTypes>, 3, 10, 1>, 3, 10>(instr, locals, device);
     break;
 case 66:
-    batch_foreach<tensor_foreach<kernel_three_body_decay_inverse<GpuTypes>, 3, 11, 1>, 3, 11>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_three_body_decay<GpuTypes>, 10, 4, 1>, 10, 4>(instr, locals, device);
     break;
 case 67:
-    batch_foreach<tensor_foreach<kernel_t_inv_min_max<GpuTypes>, 4, 2, 1>, 4, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_three_body_decay_inverse<GpuTypes>, 3, 11, 1>, 3, 11>(instr, locals, device);
     break;
 case 68:
-    batch_foreach<tensor_foreach<kernel_t_inv_value_and_min_max<GpuTypes>, 4, 3, 1>, 4, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t_inv_min_max<GpuTypes>, 4, 2, 1>, 4, 2>(instr, locals, device);
     break;
 case 69:
-    batch_foreach<tensor_foreach<kernel_t_inv_min_max_cut<GpuTypes>, 6, 2, 1>, 6, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t_inv_value_and_min_max<GpuTypes>, 4, 3, 1>, 4, 3>(instr, locals, device);
     break;
 case 70:
-    batch_foreach<tensor_foreach<kernel_t_inv_value_and_min_max_cut<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t_inv_min_max_cut<GpuTypes>, 6, 2, 1>, 6, 2>(instr, locals, device);
     break;
 case 71:
-    batch_foreach<tensor_foreach<kernel_t1_inv_min_max_doublet<GpuTypes>, 6, 2, 1>, 6, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t_inv_value_and_min_max_cut<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
     break;
 case 72:
-    batch_foreach<tensor_foreach<kernel_t1_inv_value_and_min_max_doublet<GpuTypes>, 7, 3, 1>, 7, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t1_inv_min_max_doublet<GpuTypes>, 6, 2, 1>, 6, 2>(instr, locals, device);
     break;
 case 73:
-    batch_foreach<tensor_foreach<kernel_t2_inv_min_max_doublet<GpuTypes>, 7, 2, 1>, 7, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t1_inv_value_and_min_max_doublet<GpuTypes>, 7, 3, 1>, 7, 3>(instr, locals, device);
     break;
 case 74:
-    batch_foreach<tensor_foreach<kernel_t2_inv_value_and_min_max_doublet<GpuTypes>, 8, 3, 1>, 8, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t2_inv_min_max_doublet<GpuTypes>, 7, 2, 1>, 7, 2>(instr, locals, device);
     break;
 case 75:
-    batch_foreach<tensor_foreach<kernel_s23_min_max<GpuTypes>, 6, 2, 1>, 6, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_t2_inv_value_and_min_max_doublet<GpuTypes>, 8, 3, 1>, 8, 3>(instr, locals, device);
     break;
 case 76:
-    batch_foreach<tensor_foreach<kernel_s23_value_and_min_max<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_s23_min_max<GpuTypes>, 6, 2, 1>, 6, 2>(instr, locals, device);
     break;
 case 77:
-    batch_foreach<tensor_foreach<kernel_s23_min_max_cut<GpuTypes>, 10, 2, 1>, 10, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_s23_value_and_min_max<GpuTypes>, 6, 3, 1>, 6, 3>(instr, locals, device);
     break;
 case 78:
-    batch_foreach<tensor_foreach<kernel_s23_value_and_min_max_cut<GpuTypes>, 10, 3, 1>, 10, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_s23_min_max_cut<GpuTypes>, 10, 2, 1>, 10, 2>(instr, locals, device);
     break;
 case 79:
-    batch_foreach<tensor_foreach<kernel_invariants_from_momenta<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_s23_value_and_min_max_cut<GpuTypes>, 10, 3, 1>, 10, 3>(instr, locals, device);
     break;
 case 80:
-    batch_foreach<tensor_foreach<kernel_sde2_channel_weights<GpuTypes>, 4, 1, 1>, 4, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_invariants_from_momenta<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 81:
-    batch_foreach<tensor_foreach<kernel_subchannel_weights<GpuTypes>, 6, 1, 1>, 6, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_sde2_channel_weights<GpuTypes>, 4, 1, 1>, 4, 1>(instr, locals, device);
     break;
 case 82:
-    batch_foreach<tensor_foreach<kernel_apply_subchannel_weights<GpuTypes>, 4, 1, 1>, 4, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_subchannel_weights<GpuTypes>, 6, 1, 1>, 6, 1>(instr, locals, device);
     break;
 case 83:
-    batch_foreach<tensor_foreach<kernel_pt_eta_phi_x<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_apply_subchannel_weights<GpuTypes>, 4, 1, 1>, 4, 1>(instr, locals, device);
     break;
 case 84:
-    batch_foreach<tensor_foreach<kernel_mirror_momenta<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_pt_eta_phi_x<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
     break;
 case 85:
-    batch_foreach<tensor_foreach<kernel_momenta_to_x1x2<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_mirror_momenta<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 86:
-    batch_foreach<tensor_foreach<kernel_uniform_invariant<GpuTypes>, 3, 2, 1>, 3, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_momenta_to_x1x2<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
     break;
 case 87:
-    batch_foreach<tensor_foreach<kernel_uniform_invariant_inverse<GpuTypes>, 3, 2, 1>, 3, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_uniform_invariant<GpuTypes>, 3, 2, 1>, 3, 2>(instr, locals, device);
     break;
 case 88:
-    batch_foreach<tensor_foreach<kernel_breit_wigner_invariant<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_uniform_invariant_inverse<GpuTypes>, 3, 2, 1>, 3, 2>(instr, locals, device);
     break;
 case 89:
-    batch_foreach<tensor_foreach<kernel_breit_wigner_invariant_inverse<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_breit_wigner_invariant<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
     break;
 case 90:
-    batch_foreach<tensor_foreach<kernel_stable_invariant<GpuTypes>, 4, 2, 1>, 4, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_breit_wigner_invariant_inverse<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
     break;
 case 91:
-    batch_foreach<tensor_foreach<kernel_stable_invariant_inverse<GpuTypes>, 4, 2, 1>, 4, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_stable_invariant<GpuTypes>, 4, 2, 1>, 4, 2>(instr, locals, device);
     break;
 case 92:
-    batch_foreach<tensor_foreach<kernel_stable_invariant_nu<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_stable_invariant_inverse<GpuTypes>, 4, 2, 1>, 4, 2>(instr, locals, device);
     break;
 case 93:
-    batch_foreach<tensor_foreach<kernel_stable_invariant_nu_inverse<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_stable_invariant_nu<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
     break;
 case 94:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massless<GpuTypes>, 3, 2, 1>, 3, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_stable_invariant_nu_inverse<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
     break;
 case 95:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massless_inverse<GpuTypes>, 2, 3, 1>, 2, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massless<GpuTypes>, 3, 2, 1>, 3, 2>(instr, locals, device);
     break;
 case 96:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massless_com<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massless_inverse<GpuTypes>, 2, 3, 1>, 2, 3>(instr, locals, device);
     break;
 case 97:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massive<GpuTypes>, 4, 2, 1>, 4, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massless_com<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
     break;
 case 98:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massive_inverse<GpuTypes>, 3, 3, 1>, 3, 3>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massive<GpuTypes>, 4, 2, 1>, 4, 2>(instr, locals, device);
     break;
 case 99:
-    batch_foreach<tensor_foreach<kernel_fast_rambo_massive_com<GpuTypes>, 3, 2, 1>, 3, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massive_inverse<GpuTypes>, 3, 3, 1>, 3, 3>(instr, locals, device);
     break;
 case 100:
-    batch_foreach<tensor_foreach<kernel_cut_unphysical<GpuTypes>, 4, 1, 1>, 4, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_fast_rambo_massive_com<GpuTypes>, 3, 2, 1>, 3, 2>(instr, locals, device);
     break;
 case 101:
-    batch_foreach<tensor_foreach<kernel_cut_one<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_cut_unphysical<GpuTypes>, 4, 1, 1>, 4, 1>(instr, locals, device);
     break;
 case 102:
-    batch_foreach<tensor_foreach<kernel_cut_all<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_cut_one<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
     break;
 case 103:
-    batch_foreach<tensor_foreach<kernel_cut_any<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_cut_all<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
     break;
 case 104:
-    batch_foreach<tensor_foreach<kernel_scale_transverse_energy<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_cut_any<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
     break;
 case 105:
-    batch_foreach<tensor_foreach<kernel_scale_transverse_mass<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_scale_transverse_energy<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 106:
-    batch_foreach<tensor_foreach<kernel_scale_half_transverse_mass<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_scale_transverse_mass<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 107:
-    batch_foreach<tensor_foreach<kernel_scale_partonic_energy<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_scale_half_transverse_mass<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 108:
-    batch_foreach<tensor_foreach<kernel_chili_forward<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_scale_partonic_energy<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 109:
-    batch_foreach<tensor_foreach<kernel_chili_inverse<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_chili_forward<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
     break;
 case 110:
-    op_matrix_element(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_chili_inverse<GpuTypes>, 5, 2, 1>, 5, 2>(instr, locals, device);
     break;
 case 111:
-    batch_foreach<tensor_foreach<kernel_collect_channel_weights<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
+    op_matrix_element(instr, locals, device);
     break;
 case 112:
-    batch_foreach<tensor_foreach<kernel_interpolate_pdf<GpuTypes>, 6, 1, 1>, 6, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_collect_channel_weights<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
     break;
 case 113:
-    batch_foreach<tensor_foreach<kernel_interpolate_alpha_s<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_interpolate_pdf<GpuTypes>, 6, 1, 1>, 6, 1>(instr, locals, device);
     break;
 case 114:
-    op_matmul(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_interpolate_alpha_s<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
     break;
 case 115:
-    batch_foreach<tensor_foreach_dynamic<kernel_relu<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    op_matmul(instr, locals, device);
     break;
 case 116:
-    batch_foreach<tensor_foreach_dynamic<kernel_leaky_relu<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_relu<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 117:
-    batch_foreach<tensor_foreach_dynamic<kernel_elu<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_leaky_relu<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 118:
-    batch_foreach<tensor_foreach_dynamic<kernel_gelu<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_elu<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 119:
-    batch_foreach<tensor_foreach_dynamic<kernel_sigmoid<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_gelu<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 120:
-    batch_foreach<tensor_foreach_dynamic<kernel_softplus<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_sigmoid<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 121:
-    op_rqs_reshape(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_softplus<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 122:
-    batch_foreach<tensor_foreach<kernel_rqs_find_bin<GpuTypes>, 4, 1, 2>, 4, 1>(instr, locals, device);
+    op_rqs_reshape(instr, locals, device);
     break;
 case 123:
-    batch_foreach<tensor_foreach<kernel_rqs_forward<GpuTypes>, 2, 2, 2>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_rqs_find_bin<GpuTypes>, 4, 1, 2>, 4, 1>(instr, locals, device);
     break;
 case 124:
-    batch_foreach<tensor_foreach<kernel_rqs_inverse<GpuTypes>, 2, 2, 2>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_rqs_forward<GpuTypes>, 2, 2, 2>, 2, 2>(instr, locals, device);
     break;
 case 125:
-    batch_foreach<tensor_foreach_dynamic<kernel_softmax<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_rqs_inverse<GpuTypes>, 2, 2, 2>, 2, 2>(instr, locals, device);
     break;
 case 126:
-    batch_foreach<tensor_foreach<kernel_softmax_prior<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach_dynamic<kernel_softmax<GpuTypes>, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 127:
-    batch_foreach<tensor_foreach<kernel_sample_discrete<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_softmax_prior<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 128:
-    batch_foreach<tensor_foreach<kernel_sample_discrete_inverse<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_sample_discrete<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
     break;
 case 129:
-    batch_foreach<tensor_foreach<kernel_sample_discrete_probs<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_sample_discrete_inverse<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
     break;
 case 130:
-    batch_foreach<tensor_foreach<kernel_sample_discrete_probs_inverse<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_sample_discrete_probs<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
     break;
 case 131:
-    op_discrete_histogram(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_sample_discrete_probs_inverse<GpuTypes>, 2, 2, 1>, 2, 2>(instr, locals, device);
     break;
 case 132:
-    batch_foreach<tensor_foreach<kernel_permute_momenta<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
+    op_discrete_histogram(instr, locals, device);
     break;
 case 133:
-    batch_foreach<tensor_foreach<kernel_gather<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_permute_momenta<GpuTypes>, 3, 1, 1>, 3, 1>(instr, locals, device);
     break;
 case 134:
-    batch_foreach<tensor_foreach<kernel_gather_int<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_gather<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 135:
-    batch_foreach<tensor_foreach<kernel_gather_vector<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_gather_int<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 136:
-    batch_foreach<tensor_foreach<kernel_select_int<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_gather_vector<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 137:
-    batch_foreach<tensor_foreach<kernel_select<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_select_int<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 138:
-    batch_foreach<tensor_foreach<kernel_select_vector<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_select<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 139:
-    batch_foreach<tensor_foreach<kernel_argsort<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_select_vector<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 140:
-    op_quantile(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_argsort<GpuTypes>, 1, 1, 1>, 1, 1>(instr, locals, device);
     break;
 case 141:
-    batch_foreach<tensor_foreach<kernel_one_hot<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    op_quantile(instr, locals, device);
     break;
 case 142:
-    batch_foreach<tensor_foreach<kernel_madnis_abs_weight<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_one_hot<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 143:
-    batch_foreach<tensor_foreach<kernel_madnis_softclip<GpuTypes>, 4, 1, 1>, 4, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_madnis_abs_weight<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 144:
-    batch_foreach<tensor_foreach<kernel_madnis_variance<GpuTypes>, 4, 1, 1>, 4, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_madnis_softclip<GpuTypes>, 4, 1, 1>, 4, 1>(instr, locals, device);
     break;
 case 145:
-    batch_foreach<tensor_foreach<kernel_madnis_single_channel_variance<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_madnis_variance<GpuTypes>, 4, 1, 1>, 4, 1>(instr, locals, device);
     break;
 case 146:
-    batch_foreach<tensor_foreach<kernel_madnis_multi_channel_variance<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_madnis_single_channel_variance<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 147:
-    op_nonzero(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_madnis_multi_channel_variance<GpuTypes>, 2, 1, 1>, 2, 1>(instr, locals, device);
     break;
 case 148:
-    op_batch_gather(instr, locals, device);
+    op_nonzero(instr, locals, device);
     break;
 case 149:
-    op_batch_scatter(instr, locals, device);
+    op_batch_gather(instr, locals, device);
     break;
 case 150:
-    op_random(instr, locals, device);
+    op_batch_scatter(instr, locals, device);
     break;
 case 151:
-    op_random_int(instr, locals, device);
+    op_random(instr, locals, device);
     break;
 case 152:
-    op_unweight(instr, locals, device);
+    op_random_int(instr, locals, device);
     break;
 case 153:
-    batch_foreach<tensor_foreach<kernel_vegas_forward<GpuTypes>, 2, 2, 2>, 2, 2>(instr, locals, device);
+    op_unweight(instr, locals, device);
     break;
 case 154:
-    batch_foreach<tensor_foreach<kernel_vegas_inverse<GpuTypes>, 2, 2, 2>, 2, 2>(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_vegas_forward<GpuTypes>, 2, 2, 2>, 2, 2>(instr, locals, device);
     break;
 case 155:
-    op_vegas_histogram(instr, locals, device);
+    batch_foreach<tensor_foreach<kernel_vegas_inverse<GpuTypes>, 2, 2, 2>, 2, 2>(instr, locals, device);
     break;
 case 156:
+    op_vegas_histogram(instr, locals, device);
+    break;
+case 157:
     op_histogram(instr, locals, device);
     break;
diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp
index 73fceb63da..36e375fa5e 100644
--- a/madspace/src/python/instruction_set.hpp
+++ b/madspace/src/python/instruction_set.hpp
@@ -34,6 +34,7 @@ void add_instructions(py::classh<FunctionBuilder>& fb) {
     fb.def("div", &FunctionBuilder::div, py::arg("in1"), py::arg("in2"));
     fb.def("reduce_sum", &FunctionBuilder::reduce_sum, py::arg("in1"));
     fb.def("reduce_sum_vector", &FunctionBuilder::reduce_sum_vector, py::arg("in1"));
+    fb.def("batch_reduce_sum", &FunctionBuilder::batch_reduce_sum, py::arg("in1"));
     fb.def("batch_reduce_mean", &FunctionBuilder::batch_reduce_mean, py::arg("in1"));
     fb.def("batch_reduce_mean_keepdim", &FunctionBuilder::batch_reduce_mean_keepdim, py::arg("in1"));
     fb.def("reduce_product", &FunctionBuilder::reduce_product, py::arg("in1"));
diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp
index 72c029e7cf..dd6f6da329 100644
--- a/madspace/src/python/madspace.cpp
+++ b/madspace/src/python/madspace.cpp
@@ -1035,6 +1035,7 @@ PYBIND11_MODULE(_madspace_py, m) {
                 std::size_t,
                 double,
                 double,
+                double,
                 double>(),
             py::arg("function"),
             py::arg("context"),
@@ -1043,7 +1044,8 @@ PYBIND11_MODULE(_madspace_py, m) {
             py::arg("step_count") = 0,
             py::arg("beta1") = 0.9,
             py::arg("beta2") = 0.999,
-            py::arg("eps") = 1e-8
+            py::arg("eps") = 1e-8,
+            py::arg("grad_clip_threshold") = 0.0
     )
         .def(
             "step",
@@ -1349,6 +1351,9 @@ PYBIND11_MODULE(_madspace_py, m) {
         .def_readwrite("adam_beta1", &MadnisTraining::Config::adam_beta1)
         .def_readwrite("adam_beta2", &MadnisTraining::Config::adam_beta2)
         .def_readwrite("adam_eps", &MadnisTraining::Config::adam_eps)
+        .def_readwrite(
+            "grad_clip_threshold", &MadnisTraining::Config::grad_clip_threshold
+        )
         .def_readwrite("buffer_capacity", &MadnisTraining::Config::buffer_capacity)
         .def_readwrite(
             "minimum_buffer_size", &MadnisTraining::Config::minimum_buffer_size
@@ -1383,6 +1388,13 @@ PYBIND11_MODULE(_madspace_py, m) {
         .def("active_channels", &MadnisTraining::active_channels)
         .def("active_channel_count", &MadnisTraining::active_channel_count);
 
+    py::classh<StatusFile>(m, "StatusFile")
+        .def(
+            py::init<const std::string&, double>(),
+            py::arg("file_name"),
+            py::arg("min_interval_sec") = 10.0
+        );
+
     py::classh<MultiMadnisTraining>(m, "MultiMadnisTraining")
         .def(
             py::init<
@@ -1390,12 +1402,14 @@ PYBIND11_MODULE(_madspace_py, m) {
                 ContextPtr,
                 const MadnisTraining::Config&,
                 const nested_vector2<std::shared_ptr<Integrand>>&,
-                const std::vector<std::optional<ChannelWeightNetwork>>&>(),
+                const std::vector<std::optional<ChannelWeightNetwork>>&,
+                std::shared_ptr<StatusFile>>(),
             py::arg("generator_context"),
             py::arg("optimizer_context"),
             py::arg("config"),
             py::arg("integrands"),
-            py::arg("cwnets")
+            py::arg("cwnets"),
+            py::arg("status_file") = std::shared_ptr<StatusFile>()
         )
         .def("train", &MultiMadnisTraining::train)
         .def("active_channels", &MultiMadnisTraining::active_channels);
@@ -1711,11 +1725,11 @@ PYBIND11_MODULE(_madspace_py, m) {
             py::init<
                 const std::vector<ContextPtr>&,
                 const std::vector<std::shared_ptr<ChannelEventGenerator>>&,
-                const std::string&,
+                std::shared_ptr<StatusFile>,
                 const GeneratorConfig&>(),
             py::arg("contexts"),
             py::arg("channels"),
-            py::arg("status_file") = "",
+            py::arg("status_file") = std::shared_ptr<StatusFile>(),
             py::arg_v(
                 "config",
                 EventGenerator::default_config,