diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 62626f327..4a4a01f3a 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 6114a4871..4923eaf6a 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 = ( @@ -447,7 +448,7 @@ def survey_phasespaces( def survey(self) -> None: phasespace_mode = self.run_card["phasespace"]["mode"] - if phasespace_mode == "multichannel": + if phasespace_mode in ["multichannel", "both", "auto"]: self.phasespaces = [ subproc.build_multichannel_phasespace() for subproc in self.subprocesses @@ -459,93 +460,103 @@ def survey(self) -> None: for subproc in self.subprocesses ] self.event_generator = self.survey_phasespaces(self.phasespaces) - elif phasespace_mode == "both": - kept_count = self.run_card["phasespace"]["simplified_channel_count"] - phasespaces_multi = [ - subproc.build_multichannel_phasespace() - for subproc in self.subprocesses - ] - evgen_multi = self.survey_phasespaces(phasespaces_multi) + else: + raise ValueError("Unknown phasespace mode") - phasespaces_flat = [ - subproc.build_flat_phasespace() - if len(subproc.meta["channels"]) > kept_count + 1 else - None - for subproc in self.subprocesses - ] - #evgen_flat = self.survey_phasespaces(phasespaces_flat, "flat") - - channel_status = evgen_multi.channel_status() - cross_sections = [] - index = 0 - for phasespace in phasespaces_multi: - channel_count = len(phasespace.channels) - cross_sections.append([ - status.mean - for status in channel_status[index:index + channel_count] - ]) - index += channel_count + channel_status = self.event_generator.channel_status() + chan_offset = 0 + madnis_enabled = False + for subproc, ps in zip(self.subprocesses, self.phasespaces): + mean = 0. + variance = 0. + count_opt = 0 + for status in channel_status[chan_offset:chan_offset + len(ps.channels)]: + mean += status.mean + variance += status.error**2 + count_opt += status.count_opt + rsd = (variance * count_opt)**0.5 / mean + subproc.set_madnis_auto_settings(rsd) + chan_offset += len(ps.channels) + if subproc.madnis_settings["enable"]: + madnis_enabled = True + if not ( + phasespace_mode == "both" or (phasespace_mode == "auto" and madnis_enabled) + ): + return - self.phasespaces = [ - ps_multi - if ps_flat is None else - subproc.simplify_phasespace(ps_multi, ps_flat, cross_secs) - for subproc, ps_multi, ps_flat, cross_secs in zip( - self.subprocesses, phasespaces_multi, phasespaces_flat, cross_sections - ) - ] + phasespaces_multi = self.phasespaces + cross_sections = [] + index = 0 + for phasespace in phasespaces_multi: + channel_count = len(phasespace.channels) + cross_sections.append([ + abs(status.mean) + for status in channel_status[index:index + channel_count] + ]) + index += channel_count + + self.phasespaces = [ + subproc.simplify_phasespace(ps_multi, cross_secs) + for subproc, ps_multi, cross_secs in zip( + self.subprocesses, phasespaces_multi, cross_sections + ) + ] + if any( + ps_multi is not ps_both + for ps_multi, ps_both in zip(phasespaces_multi, self.phasespaces) + ): self.event_generator = self.survey_phasespaces(self.phasespaces) - else: - raise ValueError("Unknown phasespace mode") 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() + if not any(subproc.madnis_settings["enable"] for subproc in self.subprocesses): return gen_args = self.run_card["generation"] run_args = self.run_card["run"] - config = ms.MadnisConfig() - config.verbosity = resolve_verbosity(run_args["verbosity"]) - config.learning_rate = madnis_args["lr"] - config.batches = madnis_args["train_batches"] - config.log_interval = madnis_args["log_interval"] - config.integration_history_length = madnis_args["integration_history_length"] - config.channel_dropping_interval = madnis_args["channel_dropping_interval"] - config.channel_dropping_threshold = madnis_args["channel_dropping_threshold"] - config.cpu_generator_batch_size = gen_args["cpu_batch_size"] - config.gpu_generator_batch_size = gen_args["gpu_batch_size"] - config.gpu_generator_batch_granularity = madnis_args["gpu_generator_batch_granularity"] - config.generator_target_size_factor = madnis_args["generator_target_size_factor"] - config.batch_size_offset = madnis_args["batch_size_offset"] - config.batch_size_per_channel = madnis_args["batch_size_per_channel"] - config.uniform_channel_ratio = madnis_args["uniform_channel_ratio"] - config.lr_schedule = madnis_args["lr_scheduler"] - config.adam_beta1 = madnis_args["adam_beta1"] - config.adam_beta2 = madnis_args["adam_beta2"] - config.adam_eps = madnis_args["adam_eps"] - config.buffer_capacity = madnis_args["buffer_capacity"] - config.minimum_buffer_size = madnis_args["minimum_buffer_size"] - config.buffered_steps = madnis_args["buffered_steps"] - config.buffer_unweighting_quantile = madnis_args["buffer_unweighting_quantile"] - config.fixed_cwnet_fraction = madnis_args["fixed_cwnet_fraction"] - config.softclip_threshold = madnis_args["softclip_threshold"] + verbosity = resolve_verbosity(run_args["verbosity"]) madnis_phasespaces = [] - integrands = [] - cwnets = [] + training_args = [] for subproc, phasespace in zip(self.subprocesses, self.phasespaces): + config = ms.MadnisConfig() + config.learning_rate = subproc.madnis_settings["lr"] + config.batches = subproc.madnis_settings["train_batches"] + config.log_interval = madnis_args["log_interval"] + config.integration_history_length = madnis_args["integration_history_length"] + config.channel_dropping_interval = madnis_args["channel_dropping_interval"] + config.channel_dropping_threshold = madnis_args["channel_dropping_threshold"] + config.cpu_generator_batch_size = gen_args["cpu_batch_size"] + config.gpu_generator_batch_size = gen_args["gpu_batch_size"] + config.gpu_generator_batch_granularity = madnis_args["gpu_generator_batch_granularity"] + config.generator_target_size_factor = madnis_args["generator_target_size_factor"] + config.batch_size_offset = madnis_args["batch_size_offset"] + config.batch_size_per_channel = madnis_args["batch_size_per_channel"] + config.uniform_channel_ratio = madnis_args["uniform_channel_ratio"] + config.lr_schedule = madnis_args["lr_scheduler"] + config.adam_beta1 = madnis_args["adam_beta1"] + config.adam_beta2 = madnis_args["adam_beta2"] + config.adam_eps = madnis_args["adam_eps"] + config.buffer_capacity = madnis_args["buffer_capacity"] + config.minimum_buffer_size = madnis_args["minimum_buffer_size"] + config.buffered_steps = madnis_args["buffered_steps"] + config.buffer_unweighting_quantile = madnis_args["buffer_unweighting_quantile"] + config.fixed_cwnet_fraction = madnis_args["fixed_cwnet_fraction"] + config.softclip_threshold = madnis_args["softclip_threshold"] + config.compressed_channel_weight_count = madnis_args["compressed_channel_weight_count"] phasespace = subproc.build_madnis(phasespace) madnis_phasespaces.append(phasespace) - integrands.append(subproc.build_integrands( - phasespace, - madnis_training=True, - drop_cuts_and_rescale=madnis_args["drop_zero_integrands"] - )) - cwnets.append(phasespace.cwnet) + training_args.append( + ms.TrainingArgs( + config=config, + integrands=subproc.build_integrands( + phasespace, + madnis_training=True, + drop_cuts_and_rescale=madnis_args["drop_zero_integrands"] + ), + cwnet=phasespace.cwnet, + ) + ) gen_context = self.contexts[0] opt_context = ms.Context( @@ -556,9 +567,8 @@ def train_madnis(self) -> None: madnis_training = ms.MultiMadnisTraining( generator_context=gen_context, optimizer_context=opt_context, - config=config, - integrands=integrands, - cwnets=cwnets, + training_args=training_args, + verbosity=verbosity, ) madnis_training.train() for phasespace, active_channels in zip( @@ -572,49 +582,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: @@ -1224,26 +1191,32 @@ def build_flat_phasespace(self) -> PhaseSpace: def simplify_phasespace( self, multi_phasespace: PhaseSpace, - flat_phasespace: PhaseSpace | None, cross_sections: list[float] - ) -> PhaseSpace: + ) -> PhaseSpace | None: assert multi_phasespace.mode == "multichannel" - kept_count = self.process.run_card["phasespace"]["simplified_channel_count"] - if len(multi_phasespace.channels) <= kept_count: + threshold = 1 - self.process.run_card["phasespace"]["combine_channel_threshold"] + kept_channels = [] + tot_cs = sum(cross_sections) + cum_cs = 0. + seen_active_flavors = set() + for index, (cs, chan) in sorted( + enumerate(zip(cross_sections, multi_phasespace.channels)), + key=lambda pair: pair[1][0], + reverse=True + ): + cum_cs += cs + has_unseen_flavors = False + for flavs in chan.active_flavors: + for flav in flavs: + if flav not in seen_active_flavors: + has_unseen_flavors = True + seen_active_flavors.add(flav) + if not has_unseen_flavors and cum_cs / tot_cs < threshold: + kept_channels.append(index) + if len(kept_channels) == len(cross_sections): return multi_phasespace - assert flat_phasespace is not None and flat_phasespace.mode == "flat" - #TODO: need to be careful here in the case of flavor sampling - #TODO: come up with some smarter heuristic than just channel cross section - #TODO: deal with resonances in a smart way - kept_channels = [ - index - for index, cs in sorted( - enumerate(cross_sections), key=lambda pair: pair[1], reverse=True - ) - ][:kept_count] - channels = [] channel_map = {} symfact = [] @@ -1272,6 +1245,7 @@ def simplify_phasespace( event_generator = channel.event_generator, )) + flat_phasespace = self.build_flat_phasespace() flat_channel = flat_phasespace.channels[0] channels.append(Channel( phasespace_mapping = flat_channel.phasespace_mapping, @@ -1299,6 +1273,35 @@ def simplify_phasespace( subchan_weights=multi_phasespace.subchan_weights, ) + def set_madnis_auto_settings(self, rsd: float): + madnis_args = self.process.run_card["madnis"] + n_out = len(self.meta["outgoing"]) + n_events = self.process.run_card["generation"]["events"] + is_gridpack = self.process.run_card["gridpack"]["save_gridpack"] + train_batches = madnis_args["train_batches"] + hidden_dim = min(max(int((7 * rsd) / 32) * 32 + 64, 64), 256) + flow_layers = 3 if rsd < 32 else 4 + lr = min(max((10000 - train_batches) / 8000 * 7e-4 + 3e-4, 3e-4), 1e-3) + if n_out <= 2: + enable = False + elif n_out == 3: + enable = rsd > 10. or n_events > 1000000 or is_gridpack + else: + enable = True + self.madnis_settings = { + "enable": enable, + "flow_layers": flow_layers, + "flow_hidden_dim": hidden_dim, + "discrete_hidden_dim": hidden_dim, + "cwnet_hidden_dim": hidden_dim, + "train_batches": train_batches, + "lr": lr, + } + for key, value in self.madnis_settings.items(): + run_card_value = madnis_args[key] + if run_card_value != "auto": + self.madnis_settings[key] = run_card_value + def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: madnis_args = self.process.run_card["madnis"] channels = [] @@ -1314,7 +1317,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: prefix=f"{prefix}.discrete_flow_before", dims_with_prior=[], condition_dim=0, - subnet_hidden_dim=madnis_args["discrete_hidden_dim"], + subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], subnet_layers=madnis_args["discrete_layers"], subnet_activation=self.activation(madnis_args["discrete_activation"]), ) @@ -1327,8 +1330,8 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: condition_dim=cond_dim, prefix=prefix, bin_count=madnis_args["flow_spline_bins"], - subnet_hidden_dim=madnis_args["flow_hidden_dim"], - subnet_layers=madnis_args["flow_layers"], + subnet_hidden_dim=self.madnis_settings["flow_hidden_dim"], + subnet_layers=self.madnis_settings["flow_layers"], subnet_activation=self.activation(madnis_args["flow_activation"]), invert_spline=madnis_args["flow_invert_spline"], ) @@ -1347,7 +1350,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: prefix=f"{prefix}.discrete_flow_after", dims_with_prior=[0], condition_dim=cond_dim, - subnet_hidden_dim=madnis_args["discrete_hidden_dim"], + subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], subnet_layers=madnis_args["discrete_layers"], subnet_activation=self.activation(madnis_args["discrete_activation"]), ) @@ -1415,7 +1418,7 @@ def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork: cwnet = ms.ChannelWeightNetwork( channel_count=channel_count, particle_count=self.particle_count, - hidden_dim=madnis_args["cwnet_hidden_dim"], + hidden_dim=self.madnis_settings["cwnet_hidden_dim"], layers=madnis_args["cwnet_layers"], activation=self.activation(madnis_args["cwnet_activation"]), prefix=f"subproc{self.subproc_id}.cwnet", @@ -1492,6 +1495,7 @@ def build_integrands( input_momentum_fraction=True, ) partial_weights = self.process.run_card["generation"]["systematics"] + madnis_args = self.process.run_card["madnis"] integrands = [] for channel in phasespace.channels: integrands.append(ms.Integrand( @@ -1516,6 +1520,7 @@ def build_integrands( flavor_remap, flavor_factors, flavor_mirror, + madnis_args["compressed_channel_weight_count"], )) #print(integrands[0].function()) #for i in integrands: print(i.function()) diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 28410c262..607bb37b3 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -69,7 +69,7 @@ sde_strategy = %(phasespace.sde_strategy)s #options: diagrams, denominators decays = %(phasespace.decays)s # options: all, massive, none t_channel = %(phasespace.t_channel)s # options: propagator, rambo, chili flat_mode = %(phasespace.flat_mode)s # options: propagator, rambo, chili -simplified_channel_count = %(phasespace.simplified_channel_count)s +combine_channel_threshold = %(phasespace.combine_channel_threshold)s invariant_power = %(phasespace.invariant_power)s bw_cutoff = %(phasespace.bw_cutoff)s @@ -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 @@ -136,3 +137,4 @@ batch_size_threshold = %(madnis.batch_size_threshold)s channel_grouping_mode = %(madnis.channel_grouping_mode)s # options: none, uniform, learned fixed_cwnet_fraction = %(madnis.fixed_cwnet_fraction)s softclip_threshold = %(madnis.softclip_threshold)s +compressed_channel_weight_count = %(madnis.compressed_channel_weight_count)s diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 5d5b0c63d..3096c5953 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6425,11 +6425,19 @@ def __init__(self, *args, **opts): # ------------------------------------------------------------------ # parameter declaration # ------------------------------------------------------------------ - def add_toml_param(self, section, key, value, gridpack=False, **opts): + def add_toml_param(self, section, key, value, gridpack=False, auto=False, **opts): """Declare one fixed (typed) TOML parameter belonging to ``section``. ``gridpack=True`` marks the parameter as relevant during gridpack - execution; such params are written to ``grid_run_card.toml``.""" + execution; such params are written to ``grid_run_card.toml``. + + ``auto=True`` declares a numeric parameter whose default is + determined automatically: ``value`` fixes the accepted type (int/ + float/...) and provides the internal placeholder default, but the + card reads/writes it as the string ``"auto"`` until the user sets an + explicit value of that type. This reuses the same ``auto_set`` + machinery as typing ``auto`` for any other numeric parameter, so + type-checking for genuine numeric overrides is unaffected.""" section = section.lower() key = key.lower() internal = '%s.%s' % (section, key) @@ -6442,6 +6450,8 @@ def add_toml_param(self, section, key, value, gridpack=False, **opts): self.toml_sections[section].append(key) if gridpack: self.gridpack_params.add(internal) + if auto: + self.auto_set.add(internal) def default_setup(self): """Define every parameter of the default ``run_card.toml``.""" @@ -6526,8 +6536,8 @@ def default_setup(self): self.add_toml_param('vegas', 'max_batch_size', 32000) # -------------------------- [phasespace] ---------------------- - self.add_toml_param('phasespace', 'mode', "multichannel", - allowed=['multichannel', 'flat', 'both']) + self.add_toml_param('phasespace', 'mode', "auto", + allowed=['auto', 'multichannel', 'flat', 'both']) self.add_toml_param('phasespace', 'sde_strategy', "diagrams", allowed=['diagrams', 'denominators']) self.add_toml_param('phasespace', 'decays', "all", @@ -6536,23 +6546,23 @@ def default_setup(self): allowed=['propagator', 'rambo', 'chili']) self.add_toml_param('phasespace', 'flat_mode', "rambo", allowed=['propagator', 'rambo', 'chili']) - self.add_toml_param('phasespace', 'simplified_channel_count', 10) + self.add_toml_param('phasespace', 'combine_channel_threshold', 0.01) self.add_toml_param('phasespace', 'invariant_power', 0.7) self.add_toml_param('phasespace', 'bw_cutoff', 15) # ----------------------------- [madnis] ----------------------- - self.add_toml_param('madnis', 'enable', False) - self.add_toml_param('madnis', 'flow_hidden_dim', 64) - self.add_toml_param('madnis', 'flow_layers', 3) + self.add_toml_param('madnis', 'enable', False, allowed=["auto", True, False], auto=True) + self.add_toml_param('madnis', 'flow_hidden_dim', 64, auto=True) + self.add_toml_param('madnis', 'flow_layers', 3, auto=True) self.add_toml_param('madnis', 'flow_spline_bins', 10) self.add_toml_param('madnis', 'flow_activation', "leaky_relu", allowed=['relu', 'leaky_relu', 'elu', 'gelu', 'sigmoid', 'softplus']) self.add_toml_param('madnis', 'flow_invert_spline', False) - self.add_toml_param('madnis', 'discrete_hidden_dim', 64) + self.add_toml_param('madnis', 'discrete_hidden_dim', 64, auto=True) self.add_toml_param('madnis', 'discrete_layers', 3) self.add_toml_param('madnis', 'discrete_activation', "leaky_relu", allowed=['relu', 'leaky_relu', 'elu', 'gelu', 'sigmoid', 'softplus']) - self.add_toml_param('madnis', 'cwnet_hidden_dim', 64) + self.add_toml_param('madnis', 'cwnet_hidden_dim', 64, auto=True) self.add_toml_param('madnis', 'cwnet_layers', 3) self.add_toml_param('madnis', 'cwnet_activation', "leaky_relu", allowed=['relu', 'leaky_relu', 'elu', 'gelu', 'sigmoid', 'softplus']) @@ -6564,7 +6574,7 @@ def default_setup(self): self.add_toml_param('madnis', 'batch_size_per_channel', 128) self.add_toml_param('madnis', 'generator_target_size_factor', 32) self.add_toml_param('madnis', 'gpu_generator_batch_granularity', 1000) - self.add_toml_param('madnis', 'lr', 1e-3) + self.add_toml_param('madnis', 'lr', 3e-4, auto=True) self.add_toml_param('madnis', 'lr_decay', 0.01) self.add_toml_param('madnis', 'lr_max', 3e-3) self.add_toml_param('madnis', 'lr_scheduler', "cosine", @@ -6572,10 +6582,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) @@ -6588,6 +6599,7 @@ def default_setup(self): allowed=['none', 'uniform', 'learned']) self.add_toml_param('madnis', 'fixed_cwnet_fraction', 0.33) self.add_toml_param('madnis', 'softclip_threshold', 30.0) + self.add_toml_param('madnis', 'compressed_channel_weight_count', 50) # ----------------- dynamic (free-form) sections --------------- self.dynamic_sections['multiparticles'] = collections.OrderedDict([ diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index c7a6dd13d..6a58ae617 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 360d8db14..cf6dff4cf 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]; } @@ -366,6 +370,15 @@ Value apply_subchannel_weights(Value channel_weights_in, Value subchannel_weight return instruction("apply_subchannel_weights", {channel_weights_in, subchannel_weights, channel_indices, subchannel_indices})[0]; } +std::array compress_channel_weights(Value channel_index, Value channel_weights, Value keep_count) { + auto output_vector = instruction("compress_channel_weights", {channel_index, channel_weights, keep_count}); + return {output_vector[0], output_vector[1]}; +} + +Value restore_channel_weights(Value chan_weight_values, Value chan_weight_indices, Value full_count) { + return instruction("restore_channel_weights", {chan_weight_values, chan_weight_indices, full_count})[0]; +} + Value pt_eta_phi_x(Value p_ext, Value x1, Value x2) { return instruction("pt_eta_phi_x", {p_ext, x1, x2})[0]; } diff --git a/madspace/include/madspace/compgraphs/opcode_mixin.inc b/madspace/include/madspace/compgraphs/opcode_mixin.inc index ba664afd0..748882116 100644 --- a/madspace/include/madspace/compgraphs/opcode_mixin.inc +++ b/madspace/include/madspace/compgraphs/opcode_mixin.inc @@ -19,139 +19,142 @@ 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, +compress_channel_weights = 84, +restore_channel_weights = 85, +pt_eta_phi_x = 86, +mirror_momenta = 87, +momenta_to_x1x2 = 88, +uniform_invariant = 89, +uniform_invariant_inverse = 90, +breit_wigner_invariant = 91, +breit_wigner_invariant_inverse = 92, +stable_invariant = 93, +stable_invariant_inverse = 94, +stable_invariant_nu = 95, +stable_invariant_nu_inverse = 96, +fast_rambo_massless = 97, +fast_rambo_massless_inverse = 98, +fast_rambo_massless_com = 99, +fast_rambo_massive = 100, +fast_rambo_massive_inverse = 101, +fast_rambo_massive_com = 102, +cut_unphysical = 103, +cut_one = 104, +cut_all = 105, +cut_any = 106, +scale_transverse_energy = 107, +scale_transverse_mass = 108, +scale_half_transverse_mass = 109, +scale_partonic_energy = 110, +chili_forward = 111, +chili_inverse = 112, +matrix_element = 113, +collect_channel_weights = 114, +interpolate_pdf = 115, +interpolate_alpha_s = 116, +matmul = 117, +relu = 118, +leaky_relu = 119, +elu = 120, +gelu = 121, +sigmoid = 122, +softplus = 123, +rqs_reshape = 124, +rqs_find_bin = 125, +rqs_forward = 126, +rqs_inverse = 127, +softmax = 128, +softmax_prior = 129, +sample_discrete = 130, +sample_discrete_inverse = 131, +sample_discrete_probs = 132, +sample_discrete_probs_inverse = 133, +discrete_histogram = 134, +permute_momenta = 135, +gather = 136, +gather_int = 137, +gather_vector = 138, +select_int = 139, +select = 140, +select_vector = 141, +argsort = 142, +quantile = 143, +one_hot = 144, +madnis_abs_weight = 145, +madnis_softclip = 146, +madnis_variance = 147, +madnis_single_channel_variance = 148, +madnis_multi_channel_variance = 149, +nonzero = 150, +batch_gather = 151, +batch_scatter = 152, +random = 153, +random_int = 154, +unweight = 155, +vegas_forward = 156, +vegas_inverse = 157, +vegas_histogram = 158, +histogram = 159 diff --git a/madspace/include/madspace/compgraphs/type.hpp b/madspace/include/madspace/compgraphs/type.hpp index f0d8ada0b..787e20b21 100644 --- a/madspace/include/madspace/compgraphs/type.hpp +++ b/madspace/include/madspace/compgraphs/type.hpp @@ -118,6 +118,9 @@ const Type batch_four_vec{DataType::dt_float, batch_size, {4}}; inline Type batch_float_array(int count) { return {DataType::dt_float, batch_size, {count}}; } +inline Type batch_int_array(int count) { + return {DataType::dt_int, batch_size, {count}}; +} inline Type batch_four_vec_array(int count) { return {DataType::dt_float, batch_size, {count, 4}}; } diff --git a/madspace/include/madspace/driver.hpp b/madspace/include/madspace/driver.hpp index 1982b3614..d17ce901d 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 c8df5138c..fafb83749 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 build_function_impl( + FunctionBuilder& fb, const NamedVector& 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 0e9428dff..ebc4cd3a2 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include @@ -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& contexts, const std::vector>& channels, - const std::string& status_file = "", + std::shared_ptr 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 _start_time; std::size_t _start_cpu_microsec; std::chrono::time_point _last_print_time; - std::chrono::time_point _last_status_time; PrettyBox _pretty_box_upper; PrettyBox _pretty_box_lower; - std::string _status_file; + std::shared_ptr _status_file; std::unordered_map _timing_data; bool start_jobs(); diff --git a/madspace/include/madspace/driver/madnis_training.hpp b/madspace/include/madspace/driver/madnis_training.hpp index 932760119..50460ae1f 100644 --- a/madspace/include/madspace/driver/madnis_training.hpp +++ b/madspace/include/madspace/driver/madnis_training.hpp @@ -1,11 +1,13 @@ #pragma once #include +#include #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 { @@ -17,7 +19,6 @@ class MadnisTraining { } struct Config { - Verbosity verbosity = Verbosity::log; double learning_rate = 1e-3; std::size_t batches = 1000; std::size_t log_interval = 100; @@ -35,12 +36,14 @@ 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; double buffer_unweighting_quantile = 0.99; double fixed_cwnet_fraction = 0.33; double softclip_threshold = 0.0; + std::size_t compressed_channel_weight_count = 50; }; MadnisTraining( ContextPtr generator_context, @@ -49,10 +52,35 @@ class MadnisTraining { const std::vector>& integrands, const std::optional& cwnet ); + const Config& config() const { return _config; } void train_step(std::size_t batch_index); std::vector 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& status_batches() const { return _status_batches; } + const std::vector& status_losses() const { return _status_losses; } + const std::vector& status_channel_counts() const { + return _status_channel_counts; + } + const std::vector& status_learning_rates() const { + return _status_learning_rates; + } + const std::vector& status_buffered_fractions() const { + return _status_buffered_fractions; + } + const std::vector& status_generated_events() const { + return _status_generated_events; + } + const std::vector& status_buffer_sizes() const { + return _status_buffer_sizes; + } private: struct SampleBatch { @@ -93,8 +121,12 @@ class MadnisTraining { TensorVec build_buffered_training_batch(const std::vector& counts); void process_job_results(const std::vector& job_ids); void buffer_store(ChannelData& channel, SampleBatch& samples); - void - update_history(const TensorVec& results, const std::vector& counts); + void update_history( + const TensorVec& results, + const std::vector& counts, + double learning_rate, + bool buffered + ); void drop_channels(); void freeze_cwnet(); @@ -109,7 +141,17 @@ class MadnisTraining { std::vector _channels; std::unordered_map _running_jobs; std::vector _loss_history; + std::vector _lr_history; + std::vector _buffered_history; std::size_t _loss_history_index = 0; + std::vector _status_batches; + std::vector _status_losses; + std::vector _status_channel_counts; + std::vector _status_learning_rates; + std::vector _status_buffered_fractions; + std::vector _status_generated_events; + std::vector _status_buffer_sizes; + std::size_t _generated_event_count = 0; std::size_t _job_id = 0; Tensor _generator_params; std::vector _arg_permutation; @@ -119,12 +161,18 @@ class MadnisTraining { class MultiMadnisTraining { public: + struct TrainingArgs { + MadnisTraining::Config config; + std::vector> integrands; + std::optional cwnet; + }; + MultiMadnisTraining( ContextPtr generator_context, ContextPtr optimizer_context, - const MadnisTraining::Config& config, - const nested_vector2>& integrands, - const std::vector>& cwnets + const std::vector& training_args, + Verbosity verbosity = Verbosity::log, + std::shared_ptr status_file = nullptr ); void train(); nested_vector2 active_channels() const; @@ -137,14 +185,16 @@ 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; + Verbosity _verbosity; std::vector _subprocesses; std::chrono::time_point _start_time; std::size_t _start_cpu_microsec; std::chrono::time_point _last_print_time; PrettyBox _pretty_box_upper; PrettyBox _pretty_box_lower; + std::shared_ptr _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 000000000..3cea8d175 --- /dev/null +++ b/madspace/include/madspace/driver/status_file.hpp @@ -0,0 +1,36 @@ +#pragma once + +#include +#include + +#include + +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 _min_interval; + std::chrono::time_point _last_write_time; + bool _has_written = false; + nlohmann::json _content = nlohmann::json::object(); +}; + +} // namespace madspace diff --git a/madspace/include/madspace/phasespace/integrand.hpp b/madspace/include/madspace/phasespace/integrand.hpp index 9a00bc838..384aa30c2 100644 --- a/madspace/include/madspace/phasespace/integrand.hpp +++ b/madspace/include/madspace/phasespace/integrand.hpp @@ -60,7 +60,8 @@ class Integrand : public FunctionGenerator { const nested_vector2& active_flavors = {}, const std::vector& flavor_remap = {}, const std::vector& flavor_factors = {}, - const std::vector& flavor_mirror = {} + const std::vector& flavor_mirror = {}, + std::size_t compressed_channel_weight_count = 50 ); std::size_t particle_count() const { return _mapping.particle_count(); } bool madnis_training() const { return _madnis_training; } @@ -126,6 +127,7 @@ class Integrand : public FunctionGenerator { std::optional _chan_weight_net; std::vector _chan_weight_remap; me_int_t _remapped_chan_count; + std::size_t _compressed_channel_weight_count; bool _madnis_training; bool _drop_cuts_and_rescale; bool _partial_weights; diff --git a/madspace/include/madspace/phasespace/madnis.hpp b/madspace/include/madspace/phasespace/madnis.hpp index 9b76f1907..7bb5bc356 100644 --- a/madspace/include/madspace/phasespace/madnis.hpp +++ b/madspace/include/madspace/phasespace/madnis.hpp @@ -12,7 +12,8 @@ class MadnisLoss : public FunctionGenerator { MadnisLoss( const std::vector>& functions, const std::optional& cwnet, - double softclip_threshold = 0.0 + double softclip_threshold = 0.0, + std::size_t compressed_channel_weight_count = 50 ); private: @@ -23,6 +24,7 @@ class MadnisLoss : public FunctionGenerator { std::vector> _functions; std::optional _cwnet; double _softclip_threshold; + std::size_t _compressed_channel_weight_count; }; } // namespace madspace diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index cc8055381..079c1dfff 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 @@ -1649,6 +1661,41 @@ apply_subchannel_weights: type: [float, d] desc: +compress_channel_weights: + inputs: + - name: channel_index + type: [int] + desc: + - name: channel_weights + type: [float, n] + desc: + - name: keep_count + type: [size, k] + desc: + outputs: + - name: chan_weight_values + type: [float, k] + desc: + - name: chan_weight_indices + type: [int, k] + desc: + +restore_channel_weights: + inputs: + - name: chan_weight_values + type: [float, k] + desc: + - name: chan_weight_indices + type: [int, k] + desc: + - name: full_count + type: [size, n] + desc: + outputs: + - name: channel_weights + type: [float, n] + desc: + pt_eta_phi_x: inputs: - name: p_ext diff --git a/madspace/madspace/_madspace_py.pyi b/madspace/madspace/_madspace_py.pyi index d64a398e7..0faed1b22 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 31f5dbab3..000000000 --- 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 d481c7d6d..000000000 --- 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 76ee97013..000000000 --- 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 120aad48d..000000000 --- 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 41fe917e1..000000000 --- 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 c7b1729e4..000000000 --- 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 b36c82f77..000000000 --- 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 9daaaa3c2..000000000 --- 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 a2dbe6cd8..820f582f6 100644 --- a/madspace/src/compgraphs/instruction_set_mixin.inc +++ b/madspace/src/compgraphs/instruction_set_mixin.inc @@ -36,140 +36,143 @@ 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("compress_channel_weights", 84, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, true, {"k"}, true}}, {{DataType::dt_float, false, {"k"}, false}, {DataType::dt_int, false, {"k"}, false}}), + mi("restore_channel_weights", 85, true, {{DataType::dt_float, false, {"k"}, false}, {DataType::dt_int, false, {"k"}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("pt_eta_phi_x", 86, 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", 87, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), + mi("momenta_to_x1x2", 88, 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", 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}}), + mi("uniform_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}}), + mi("breit_wigner_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}, {DataType::dt_float, false, {}, false}}), + mi("breit_wigner_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}, {DataType::dt_float, false, {}, false}}), + mi("stable_invariant", 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}}), + mi("stable_invariant_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}}), + mi("stable_invariant_nu", 95, 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", 96, 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", 97, 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", 98, 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", 99, 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", 100, 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", 101, 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", 102, 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", 103, 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", 104, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("cut_all", 105, 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", 106, 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", 107, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_transverse_mass", 108, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_half_transverse_mass", 109, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_partonic_energy", 110, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("chili_forward", 111, 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", 112, 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(113, true)), + mi("collect_channel_weights", 114, 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", 115, 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", 116, 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", 117, 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", 118, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("leaky_relu", 119, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("elu", 120, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("gelu", 121, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("sigmoid", 122, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softplus", 123, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + InstructionOwner(new RqsReshapeInstruction(124, true)), + mi("rqs_find_bin", 125, 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", 126, 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", 127, 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", 128, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softmax_prior", 129, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("sample_discrete", 130, 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", 131, 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", 132, 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", 133, 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", 134, 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", 135, 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", 136, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("gather_int", 137, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), + mi("gather_vector", 138, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n", "m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("select_int", 139, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), + mi("select", 140, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("select_vector", 141, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), + mi("argsort", 142, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), + mi("quantile", 143, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, true, {}, false}}), + mi("one_hot", 144, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("madnis_abs_weight", 145, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_softclip", 146, 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", 147, 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", 148, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_multi_channel_variance", 149, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), + InstructionOwner(new NonzeroInstruction(150, true)), + InstructionOwner(new BatchGatherInstruction(151, true)), + InstructionOwner(new BatchScatterInstruction(152, true)), + InstructionOwner(new RandomInstruction(153, true)), + InstructionOwner(new RandomIntInstruction(154, true)), + InstructionOwner(new UnweightInstruction(155, true)), + mi("vegas_forward", 156, 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", 157, 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", 158, 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", 159, 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 02ed5a5be..9f0fb2bde 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -408,11 +408,12 @@ void op_batch_scatter( } template -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(0); auto output_view_flat = output.flat_view(0); - device.submit([keepdim, input_view_flat, output_view_flat, batch_size]() mutable { - TensorView input_view(input_view_flat); - TensorView 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 input_view(input_view_flat); + TensorView 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 -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(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 input_grad_view(input_grad_view_flat); TensorView 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 +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 +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 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 @@ -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 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 @@ -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 diff --git a/madspace/src/cpu/runtime_backward_mixin.inc b/madspace/src/cpu/runtime_backward_mixin.inc index 8a75bf498..b47a9ae2f 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, backward_kernel_div, 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, backward_kernel_reduce_product, 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, backward_kernel_sqrt, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); + backward_batch_foreach, backward_kernel_reduce_product, 2, 1, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; case 25: + backward_batch_foreach, backward_kernel_sqrt, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); + break; +case 26: backward_batch_foreach, backward_kernel_square, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 114: +case 117: backward_op_matmul(instr, locals, local_grads, device); break; -case 115: +case 118: backward_batch_foreach, backward_kernel_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 116: +case 119: backward_batch_foreach, backward_kernel_leaky_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 117: +case 120: backward_batch_foreach, backward_kernel_elu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 118: +case 121: backward_batch_foreach, backward_kernel_gelu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 119: +case 122: backward_batch_foreach, backward_kernel_sigmoid, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 120: +case 123: backward_batch_foreach, backward_kernel_softplus, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 121: +case 124: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 122: +case 125: backward_batch_foreach, backward_kernel_rqs_find_bin, 5, 4, 2, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 123: +case 126: backward_batch_foreach, backward_kernel_rqs_forward, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 124: +case 127: backward_batch_foreach, backward_kernel_rqs_inverse, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 125: +case 128: backward_batch_foreach, backward_kernel_softmax, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 126: +case 129: backward_batch_foreach, backward_kernel_softmax_prior, 2, 2, 1, DeviceType>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 130: +case 133: backward_batch_foreach, backward_kernel_sample_discrete_probs_inverse, 4, 2, 1, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 133: +case 136: backward_batch_foreach, backward_kernel_gather, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 137: +case 140: backward_batch_foreach, backward_kernel_select, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 142: +case 145: backward_batch_foreach, backward_kernel_madnis_abs_weight, 3, 2, 1, DeviceType>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 143: +case 146: backward_batch_foreach, backward_kernel_madnis_softclip, 5, 4, 1, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 144: +case 147: backward_batch_foreach, backward_kernel_madnis_variance, 5, 4, 1, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 145: +case 148: backward_batch_foreach, backward_kernel_madnis_single_channel_variance, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 146: +case 149: backward_batch_foreach, backward_kernel_madnis_multi_channel_variance, 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 64a3974be..7e6f47ae2 100644 --- a/madspace/src/cpu/runtime_mixin.inc +++ b/madspace/src/cpu/runtime_mixin.inc @@ -65,410 +65,419 @@ case 20: batch_foreach, kernel_reduce_sum_vector, 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, kernel_reduce_product, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_batch_reduce_mean_keepdim(instr, locals, device); break; case 24: - batch_foreach, kernel_sqrt, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_reduce_product, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 25: - batch_foreach, kernel_square, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_sqrt, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 26: - batch_foreach, kernel_min, 2, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_square, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 27: - batch_foreach, kernel_max, 2, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_min, 2, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 28: - batch_foreach, kernel_obs_sqrt_s, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_max, 2, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 29: - batch_foreach, kernel_obs_e, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_sqrt_s, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 30: - batch_foreach, kernel_obs_px, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_e, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 31: - batch_foreach, kernel_obs_py, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_px, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 32: - batch_foreach, kernel_obs_pz, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_py, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 33: - batch_foreach, kernel_obs_mass, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_pz, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 34: - batch_foreach, kernel_obs_pt, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_mass, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 35: - batch_foreach, kernel_obs_p_mag, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_pt, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 36: - batch_foreach, kernel_obs_phi, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_p_mag, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 37: - batch_foreach, kernel_obs_theta, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_phi, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 38: - batch_foreach, kernel_obs_y, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_theta, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 39: - batch_foreach, kernel_obs_y_abs, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_y, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 40: - batch_foreach, kernel_obs_eta, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_y_abs, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 41: - batch_foreach, kernel_obs_eta_abs, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_obs_eta, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 42: - batch_foreach, kernel_obs_delta_eta, 2, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_obs_eta_abs, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 43: - batch_foreach, kernel_obs_delta_phi, 2, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_obs_delta_eta, 2, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 44: - batch_foreach, kernel_obs_delta_r, 2, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_obs_delta_phi, 2, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 45: - batch_foreach, kernel_boost_beam, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_obs_delta_r, 2, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 46: - batch_foreach, kernel_boost_beam_inverse, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_boost_beam, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 47: - batch_foreach, kernel_com_p_in, 1, 2, 1, DeviceType>, 1, 2>(instr, locals, device); + batch_foreach, kernel_boost_beam_inverse, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 48: - batch_foreach, kernel_r_to_x1x2, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device); + batch_foreach, kernel_com_p_in, 1, 2, 1, DeviceType>, 1, 2>(instr, locals, device); break; case 49: - batch_foreach, kernel_x1x2_to_r, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_r_to_x1x2, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device); break; case 50: - batch_foreach, kernel_diff_cross_section, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_x1x2_to_r, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 51: - batch_foreach, kernel_two_body_decay_com, 5, 3, 1, DeviceType>, 5, 3>(instr, locals, device); + batch_foreach, kernel_diff_cross_section, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 52: - batch_foreach, kernel_two_body_decay_com_inverse, 2, 6, 1, DeviceType>, 2, 6>(instr, locals, device); + batch_foreach, kernel_two_body_decay_com, 5, 3, 1, DeviceType>, 5, 3>(instr, locals, device); break; case 53: - batch_foreach, kernel_two_body_decay, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); + batch_foreach, kernel_two_body_decay_com_inverse, 2, 6, 1, DeviceType>, 2, 6>(instr, locals, device); break; case 54: - batch_foreach, kernel_two_body_decay_inverse, 2, 7, 1, DeviceType>, 2, 7>(instr, locals, device); + batch_foreach, kernel_two_body_decay, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 55: - batch_foreach, kernel_two_to_two_particle_scattering_com, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); + batch_foreach, kernel_two_body_decay_inverse, 2, 7, 1, DeviceType>, 2, 7>(instr, locals, device); break; case 56: - batch_foreach, kernel_two_to_two_particle_scattering_com_inverse, 4, 4, 1, DeviceType>, 4, 4>(instr, locals, device); + batch_foreach, kernel_two_to_two_particle_scattering_com, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 57: - batch_foreach, kernel_two_to_two_particle_scattering, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); + batch_foreach, kernel_two_to_two_particle_scattering_com_inverse, 4, 4, 1, DeviceType>, 4, 4>(instr, locals, device); break; case 58: - batch_foreach, kernel_two_to_two_particle_scattering_inverse, 4, 4, 1, DeviceType>, 4, 4>(instr, locals, device); + batch_foreach, kernel_two_to_two_particle_scattering, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 59: - batch_foreach, kernel_two_to_three_particle_scattering, 8, 3, 1, DeviceType>, 8, 3>(instr, locals, device); + batch_foreach, kernel_two_to_two_particle_scattering_inverse, 4, 4, 1, DeviceType>, 4, 4>(instr, locals, device); break; case 60: - batch_foreach, kernel_two_to_three_particle_scattering_inverse, 7, 4, 1, DeviceType>, 7, 4>(instr, locals, device); + batch_foreach, kernel_two_to_three_particle_scattering, 8, 3, 1, DeviceType>, 8, 3>(instr, locals, device); break; case 61: - batch_foreach, kernel_double_t_scattering, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); + batch_foreach, kernel_two_to_three_particle_scattering_inverse, 7, 4, 1, DeviceType>, 7, 4>(instr, locals, device); break; case 62: - batch_foreach, kernel_double_t_scattering_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_double_t_scattering, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 63: - batch_foreach, kernel_three_body_decay_com, 9, 4, 1, DeviceType>, 9, 4>(instr, locals, device); + batch_foreach, kernel_double_t_scattering_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 64: - batch_foreach, kernel_three_body_decay_com_inverse, 3, 10, 1, DeviceType>, 3, 10>(instr, locals, device); + batch_foreach, kernel_three_body_decay_com, 9, 4, 1, DeviceType>, 9, 4>(instr, locals, device); break; case 65: - batch_foreach, kernel_three_body_decay, 10, 4, 1, DeviceType>, 10, 4>(instr, locals, device); + batch_foreach, kernel_three_body_decay_com_inverse, 3, 10, 1, DeviceType>, 3, 10>(instr, locals, device); break; case 66: - batch_foreach, kernel_three_body_decay_inverse, 3, 11, 1, DeviceType>, 3, 11>(instr, locals, device); + batch_foreach, kernel_three_body_decay, 10, 4, 1, DeviceType>, 10, 4>(instr, locals, device); break; case 67: - batch_foreach, kernel_t_inv_min_max, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_three_body_decay_inverse, 3, 11, 1, DeviceType>, 3, 11>(instr, locals, device); break; case 68: - batch_foreach, kernel_t_inv_value_and_min_max, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); + batch_foreach, kernel_t_inv_min_max, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 69: - batch_foreach, kernel_t_inv_min_max_cut, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); + batch_foreach, kernel_t_inv_value_and_min_max, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); break; case 70: - batch_foreach, kernel_t_inv_value_and_min_max_cut, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); + batch_foreach, kernel_t_inv_min_max_cut, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); break; case 71: - batch_foreach, kernel_t1_inv_min_max_doublet, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); + batch_foreach, kernel_t_inv_value_and_min_max_cut, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 72: - batch_foreach, kernel_t1_inv_value_and_min_max_doublet, 7, 3, 1, DeviceType>, 7, 3>(instr, locals, device); + batch_foreach, kernel_t1_inv_min_max_doublet, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); break; case 73: - batch_foreach, kernel_t2_inv_min_max_doublet, 7, 2, 1, DeviceType>, 7, 2>(instr, locals, device); + batch_foreach, kernel_t1_inv_value_and_min_max_doublet, 7, 3, 1, DeviceType>, 7, 3>(instr, locals, device); break; case 74: - batch_foreach, kernel_t2_inv_value_and_min_max_doublet, 8, 3, 1, DeviceType>, 8, 3>(instr, locals, device); + batch_foreach, kernel_t2_inv_min_max_doublet, 7, 2, 1, DeviceType>, 7, 2>(instr, locals, device); break; case 75: - batch_foreach, kernel_s23_min_max, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); + batch_foreach, kernel_t2_inv_value_and_min_max_doublet, 8, 3, 1, DeviceType>, 8, 3>(instr, locals, device); break; case 76: - batch_foreach, kernel_s23_value_and_min_max, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); + batch_foreach, kernel_s23_min_max, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); break; case 77: - batch_foreach, kernel_s23_min_max_cut, 10, 2, 1, DeviceType>, 10, 2>(instr, locals, device); + batch_foreach, kernel_s23_value_and_min_max, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 78: - batch_foreach, kernel_s23_value_and_min_max_cut, 10, 3, 1, DeviceType>, 10, 3>(instr, locals, device); + batch_foreach, kernel_s23_min_max_cut, 10, 2, 1, DeviceType>, 10, 2>(instr, locals, device); break; case 79: - batch_foreach, kernel_invariants_from_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_s23_value_and_min_max_cut, 10, 3, 1, DeviceType>, 10, 3>(instr, locals, device); break; case 80: - batch_foreach, kernel_sde2_channel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_invariants_from_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 81: - batch_foreach, kernel_subchannel_weights, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_sde2_channel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 82: - batch_foreach, kernel_apply_subchannel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_subchannel_weights, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 83: - batch_foreach, kernel_pt_eta_phi_x, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_apply_subchannel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 84: - batch_foreach, kernel_mirror_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_compress_channel_weights, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 85: - batch_foreach, kernel_momenta_to_x1x2, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_restore_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 86: - batch_foreach, kernel_uniform_invariant, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_pt_eta_phi_x, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 87: - batch_foreach, kernel_uniform_invariant_inverse, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_mirror_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 88: - batch_foreach, kernel_breit_wigner_invariant, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_momenta_to_x1x2, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 89: - batch_foreach, kernel_breit_wigner_invariant_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_uniform_invariant, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 90: - batch_foreach, kernel_stable_invariant, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_uniform_invariant_inverse, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 91: - batch_foreach, kernel_stable_invariant_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_breit_wigner_invariant, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 92: - batch_foreach, kernel_stable_invariant_nu, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_breit_wigner_invariant_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 93: - batch_foreach, kernel_stable_invariant_nu_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 94: - batch_foreach, kernel_fast_rambo_massless, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 95: - batch_foreach, kernel_fast_rambo_massless_inverse, 2, 3, 1, DeviceType>, 2, 3>(instr, locals, device); + batch_foreach, kernel_stable_invariant_nu, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 96: - batch_foreach, kernel_fast_rambo_massless_com, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant_nu_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 97: - batch_foreach, kernel_fast_rambo_massive, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 98: - batch_foreach, kernel_fast_rambo_massive_inverse, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless_inverse, 2, 3, 1, DeviceType>, 2, 3>(instr, locals, device); break; case 99: - batch_foreach, kernel_fast_rambo_massive_com, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless_com, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 100: - batch_foreach, kernel_cut_unphysical, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 101: - batch_foreach, kernel_cut_one, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive_inverse, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device); break; case 102: - batch_foreach, kernel_cut_all, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive_com, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 103: - batch_foreach, kernel_cut_any, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_cut_unphysical, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 104: - batch_foreach, kernel_scale_transverse_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_cut_one, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 105: - batch_foreach, kernel_scale_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_cut_all, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 106: - batch_foreach, kernel_scale_half_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_cut_any, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 107: - batch_foreach, kernel_scale_partonic_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_scale_transverse_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 108: - batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_scale_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 109: - batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_scale_half_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 110: - op_matrix_element(instr, locals, device); + batch_foreach, kernel_scale_partonic_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 111: - batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 112: - batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 113: - batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 114: - op_matmul(instr, locals, device); + batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 115: - batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 116: - batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 117: - batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 118: - batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 119: - batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 120: - batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 121: - op_rqs_reshape(instr, locals, device); + batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 122: - batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 123: - batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 124: - batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 125: - batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); break; case 126: - batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 127: - batch_foreach, kernel_sample_discrete, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 128: - batch_foreach, kernel_sample_discrete_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 129: - batch_foreach, kernel_sample_discrete_probs, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 130: - batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_sample_discrete, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 131: - op_discrete_histogram(instr, locals, device); + batch_foreach, kernel_sample_discrete_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 132: - batch_foreach, kernel_permute_momenta, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 133: - batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 134: - batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 135: - batch_foreach, kernel_gather_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_permute_momenta, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 136: - batch_foreach, kernel_select_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 137: - batch_foreach, kernel_select, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 138: - batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 139: - batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_select_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 140: - op_quantile(instr, locals, device); + batch_foreach, kernel_select, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 141: - batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 142: - batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 143: - batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + op_quantile(instr, locals, device); break; case 144: - batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 145: - batch_foreach, kernel_madnis_single_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 146: - batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 147: - op_nonzero(instr, locals, device); + batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 148: - op_batch_gather(instr, locals, device); + batch_foreach, kernel_madnis_single_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 149: - op_batch_scatter(instr, locals, device); + batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 150: - op_random(instr, locals, device); + op_nonzero(instr, locals, device); break; case 151: - op_random_int(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 152: - op_unweight(instr, locals, device); + op_batch_scatter(instr, locals, device); break; case 153: - batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_random(instr, locals, device); break; case 154: - batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_random_int(instr, locals, device); break; case 155: - op_vegas_histogram(instr, locals, device); + op_unweight(instr, locals, device); break; case 156: + batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + break; +case 157: + batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + break; +case 158: + op_vegas_histogram(instr, locals, device); + break; +case 159: op_histogram(instr, locals, device); break; diff --git a/madspace/src/driver/adam_optimizer.cpp b/madspace/src/driver/adam_optimizer.cpp index e7959cf69..c3686e8f7 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 GradientClipper::build_function_impl( + FunctionBuilder& fb, const NamedVector& 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 a5e4760a5..dcf99ae52 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 -#include #include #include @@ -15,7 +14,7 @@ const GeneratorConfig EventGenerator::default_config = {}; EventGenerator::EventGenerator( const std::vector& contexts, const std::vector>& channels, - const std::string& status_file, + std::shared_ptr 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 c94d84fb9..9a97c2b03 100644 --- a/madspace/src/driver/madnis_training.cpp +++ b/madspace/src/driver/madnis_training.cpp @@ -27,7 +27,12 @@ MadnisTraining::MadnisTraining( _arg_permutation.push_back(integ_args.at("adaptive_prob")); if (cwnet) { _arg_permutation.push_back(integ_args.at("cwnet_input")); - _arg_permutation.push_back(integ_args.at("channel_weights")); + if (integ_args.contains("channel_weight_values")) { + _arg_permutation.push_back(integ_args.at("channel_weight_values")); + _arg_permutation.push_back(integ_args.at("channel_weight_indices")); + } else { + _arg_permutation.push_back(integ_args.at("channel_weights")); + } _arg_permutation.push_back(integ_args.at("channel_index")); } for (auto key : channel.integrand_prob->arg_types().keys()) { @@ -53,11 +58,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 +74,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 job_ids; @@ -111,6 +128,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(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> functions; functions.reserve(_channels.size()); @@ -120,7 +167,13 @@ void MadnisTraining::build_runtimes_and_optimizer() { channel.generator_runtime.reset(); } Function madnis_func = - MadnisLoss(functions, _cwnet, _config.softclip_threshold).function(); + MadnisLoss( + functions, + _cwnet, + _config.softclip_threshold, + _config.compressed_channel_weight_count + ) + .function(); if (_optimizer) { _optimizer->replace_function(madnis_func); } else { @@ -132,7 +185,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 +565,7 @@ void MadnisTraining::process_job_results(const std::vector& 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 +580,7 @@ void MadnisTraining::process_job_results(const std::vector& 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 +645,26 @@ void MadnisTraining::buffer_store(ChannelData& channel, SampleBatch& samples) { } void MadnisTraining::update_history( - const TensorVec& results, const std::vector& counts + const TensorVec& results, + const std::vector& 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()[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; @@ -717,15 +783,19 @@ void MadnisTraining::freeze_cwnet() { MultiMadnisTraining::MultiMadnisTraining( ContextPtr generator_context, ContextPtr optimizer_context, - const MadnisTraining::Config& config, - const nested_vector2>& integrands, - const std::vector>& cwnets + const std::vector& training_args, + Verbosity verbosity, + std::shared_ptr status_file ) : - _config(config) { - _subprocesses.reserve(integrands.size()); - for (auto [integ, cwnet] : zip(integrands, cwnets)) { + _verbosity(verbosity), _status_file(status_file) { + _subprocesses.reserve(training_args.size()); + for (auto args : training_args) { _subprocesses.emplace_back( - generator_context, optimizer_context, config, integ, cwnet + generator_context, + optimizer_context, + args.config, + args.integrands, + args.cwnet ); } } @@ -735,12 +805,15 @@ void MultiMadnisTraining::train() { _start_time = std::chrono::steady_clock::now(); _start_cpu_microsec = cpu_time_microsec(); for (std::size_t subproc_index = 0; auto& subproc : _subprocesses) { - for (std::size_t batch_index = 0; batch_index < _config.batches; + for (std::size_t batch_index = 0; batch_index < subproc.config().batches; ++batch_index) { subproc.train_step(batch_index); 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 == subproc.config().batches; + write_status(subproc_index, batch_index, done); } ++subproc_index; } @@ -757,7 +830,7 @@ nested_vector2 MultiMadnisTraining::active_channels() const { void MultiMadnisTraining::print_progress_init() { _last_print_time = std::chrono::steady_clock::now(); - if (_config.verbosity != Verbosity::pretty) { + if (_verbosity != Verbosity::pretty) { Logger::info("training started"); return; } @@ -785,8 +858,14 @@ void MultiMadnisTraining::print_progress_update( std::size_t chan_count ) { using namespace std::chrono_literals; - if (_config.verbosity == Verbosity::log) { - if ((batch_index + 1) % _config.log_interval != 0) { + std::size_t subproc_count = _subprocesses.size(); + std::size_t subproc_batches = _subprocesses.at(subproc_index).config().batches; + bool is_last_batch = batch_index + 1 == subproc_batches; + bool is_done = is_last_batch && subproc_index == subproc_count - 1; + + if (_verbosity == Verbosity::log) { + if ((batch_index + 1) % _subprocesses.at(subproc_index).config().log_interval != + 0) { return; } auto now = std::chrono::steady_clock::now(); @@ -795,15 +874,15 @@ void MultiMadnisTraining::print_progress_update( "training, subproc: {} / {}, batch: {} / {}, loss: {:.4f}, channels: " "{}, time: {:%H:%M:%S}", subproc_index + 1, - _subprocesses.size(), + subproc_count, batch_index + 1, - _config.batches, + subproc_batches, loss, chan_count, std::chrono::round(now - _start_time) ) ); - if (batch_index + 1 == _config.batches) { + if (is_done) { double wall_time_sec = (now - _start_time) / 1.0s; double cpu_time_sec = (cpu_time_microsec() - _start_cpu_microsec) / 1e6; Logger::info( @@ -812,17 +891,13 @@ void MultiMadnisTraining::print_progress_update( ) ); } - } else if (_config.verbosity == Verbosity::pretty) { + } else if (_verbosity == Verbosity::pretty) { auto now = std::chrono::steady_clock::now(); - if (now - _last_print_time < 0.1s && batch_index + 1 != _config.batches) { + if (now - _last_print_time < 0.1s && !is_last_batch) { return; } _last_print_time = now; - std::size_t subproc_count = _subprocesses.size(); - bool is_last_batch = batch_index + 1 == _config.batches; - bool is_done = is_last_batch && subproc_index == subproc_count - 1; - std::string time_str; if (is_done) { double wall_time_sec = (now - _start_time) / 1.0s; @@ -835,13 +910,13 @@ void MultiMadnisTraining::print_progress_update( ); } std::string batch_str = - std::format("{} / {}", batch_index + 1, _config.batches); + std::format("{} / {}", batch_index + 1, subproc_batches); if (subproc_count == 1) { std::string progress_bar = is_done ? "" : format_progress( - static_cast(batch_index + 1) / _config.batches, 52 + static_cast(batch_index + 1) / subproc_batches, 52 ); _pretty_box_upper.set_column( 1, @@ -852,23 +927,34 @@ void MultiMadnisTraining::print_progress_update( ); _pretty_box_upper.print_update(); } else { + // Weight each subprocess's contribution to the global progress bar by its + // own batch count, since subprocesses no longer share a single batch count. + std::size_t total_batches = 0; + std::size_t batches_before = 0; + for (std::size_t i = 0; i < subproc_count; ++i) { + std::size_t batches = _subprocesses.at(i).config().batches; + total_batches += batches; + if (i < subproc_index) { + batches_before += batches; + } + } + std::string progress_bar, progress_bar_all; std::string subproc_str = - std::format("{} / {}", subproc_index, subproc_count); + std::format("{} / {}", subproc_index + 1, subproc_count); if (!is_last_batch) { progress_bar = format_progress( - static_cast(batch_index + 1) / _config.batches, 34 + static_cast(batch_index + 1) / subproc_batches, 34 ); progress_bar_all = format_progress( - static_cast( - subproc_index * _config.batches + batch_index + 1 - ) / (subproc_count * _config.batches), + static_cast(batches_before + batch_index + 1) / + total_batches, 52 ); } else if (!is_done) { progress_bar_all = format_progress( - static_cast((subproc_index + 1) * _config.batches + 1) / - (subproc_count * _config.batches), + static_cast(batches_before + subproc_batches + 1) / + total_batches, 52 ); } @@ -887,3 +973,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) % _subprocesses.at(0).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", subproc.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 000000000..8ece6ed93 --- /dev/null +++ b/madspace/src/driver/status_file.cpp @@ -0,0 +1,44 @@ +#include "madspace/driver/status_file.hpp" + +#include +#include +#include + +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 87ff0b0be..af9dde25d 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(), - out.view() - ); + if (is_mean) { + launch_kernel( + kernel_div_batch_size, + 1, + device.stream(), + batch_size, + out.view(), + out.view() + ); + } 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(), - grad.view() - ); temp.reset(device); output_grad.reset(device); + if (is_mean) { + launch_kernel( + kernel_div_batch_size, + 1, + device.stream(), + batch_size, + grad.view(), + grad.view() + ); + } } 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(), - grad.view() - ); + if (is_mean) { + launch_kernel( + kernel_div_batch_size, + 1, + device.stream(), + input.size(0), + output_grad.view(), + grad.view() + ); + } 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 452946c95..00b3547c5 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, 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, 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, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); + backward_batch_foreach, 2, 1, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; case 25: + backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); + break; +case 26: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 114: +case 117: backward_op_matmul(instr, locals, local_grads, device); break; -case 115: +case 118: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 116: +case 119: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 117: +case 120: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 118: +case 121: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 119: +case 122: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 120: +case 123: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 121: +case 124: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 122: +case 125: backward_batch_foreach, 5, 4, 2>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 123: +case 126: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 124: +case 127: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 125: +case 128: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 126: +case 129: backward_batch_foreach, 2, 2, 1>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 130: +case 133: backward_batch_foreach, 4, 2, 1>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 133: +case 136: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 137: +case 140: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 142: +case 145: backward_batch_foreach, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 143: +case 146: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 144: +case 147: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 145: +case 148: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 146: +case 149: backward_batch_foreach, 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 316c746a3..4ce0dc8f5 100644 --- a/madspace/src/gpu/runtime_mixin.inc +++ b/madspace/src/gpu/runtime_mixin.inc @@ -65,410 +65,419 @@ case 20: batch_foreach, 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, 1, 1, 1>, 1, 1>(instr, locals, device); + op_batch_reduce_mean_keepdim(instr, locals, device); break; case 24: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 25: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 26: - batch_foreach, 2, 1>, 2, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 27: - batch_foreach, 2, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1>, 2, 1>(instr, locals, device); break; case 28: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 2, 1>, 2, 1>(instr, locals, device); break; case 29: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 30: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 31: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 32: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 33: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 34: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 35: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 36: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 37: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 38: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 39: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 40: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 41: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 42: - batch_foreach, 2, 1>, 2, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 43: - batch_foreach, 2, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1>, 2, 1>(instr, locals, device); break; case 44: - batch_foreach, 2, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1>, 2, 1>(instr, locals, device); break; case 45: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 2, 1>, 2, 1>(instr, locals, device); break; case 46: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 47: - batch_foreach, 1, 2, 1>, 1, 2>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 48: - batch_foreach, 3, 3, 1>, 3, 3>(instr, locals, device); + batch_foreach, 1, 2, 1>, 1, 2>(instr, locals, device); break; case 49: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 3, 3, 1>, 3, 3>(instr, locals, device); break; case 50: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 51: - batch_foreach, 5, 3, 1>, 5, 3>(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 52: - batch_foreach, 2, 6, 1>, 2, 6>(instr, locals, device); + batch_foreach, 5, 3, 1>, 5, 3>(instr, locals, device); break; case 53: - batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); + batch_foreach, 2, 6, 1>, 2, 6>(instr, locals, device); break; case 54: - batch_foreach, 2, 7, 1>, 2, 7>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 55: - batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); + batch_foreach, 2, 7, 1>, 2, 7>(instr, locals, device); break; case 56: - batch_foreach, 4, 4, 1>, 4, 4>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 57: - batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); + batch_foreach, 4, 4, 1>, 4, 4>(instr, locals, device); break; case 58: - batch_foreach, 4, 4, 1>, 4, 4>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 59: - batch_foreach, 8, 3, 1>, 8, 3>(instr, locals, device); + batch_foreach, 4, 4, 1>, 4, 4>(instr, locals, device); break; case 60: - batch_foreach, 7, 4, 1>, 7, 4>(instr, locals, device); + batch_foreach, 8, 3, 1>, 8, 3>(instr, locals, device); break; case 61: - batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); + batch_foreach, 7, 4, 1>, 7, 4>(instr, locals, device); break; case 62: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 63: - batch_foreach, 9, 4, 1>, 9, 4>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 64: - batch_foreach, 3, 10, 1>, 3, 10>(instr, locals, device); + batch_foreach, 9, 4, 1>, 9, 4>(instr, locals, device); break; case 65: - batch_foreach, 10, 4, 1>, 10, 4>(instr, locals, device); + batch_foreach, 3, 10, 1>, 3, 10>(instr, locals, device); break; case 66: - batch_foreach, 3, 11, 1>, 3, 11>(instr, locals, device); + batch_foreach, 10, 4, 1>, 10, 4>(instr, locals, device); break; case 67: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 3, 11, 1>, 3, 11>(instr, locals, device); break; case 68: - batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 69: - batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); + batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); break; case 70: - batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); + batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); break; case 71: - batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 72: - batch_foreach, 7, 3, 1>, 7, 3>(instr, locals, device); + batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); break; case 73: - batch_foreach, 7, 2, 1>, 7, 2>(instr, locals, device); + batch_foreach, 7, 3, 1>, 7, 3>(instr, locals, device); break; case 74: - batch_foreach, 8, 3, 1>, 8, 3>(instr, locals, device); + batch_foreach, 7, 2, 1>, 7, 2>(instr, locals, device); break; case 75: - batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); + batch_foreach, 8, 3, 1>, 8, 3>(instr, locals, device); break; case 76: - batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); + batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); break; case 77: - batch_foreach, 10, 2, 1>, 10, 2>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 78: - batch_foreach, 10, 3, 1>, 10, 3>(instr, locals, device); + batch_foreach, 10, 2, 1>, 10, 2>(instr, locals, device); break; case 79: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 10, 3, 1>, 10, 3>(instr, locals, device); break; case 80: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 81: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 82: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 83: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 84: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 85: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 86: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 87: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 88: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 89: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 90: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 91: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 92: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 93: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 94: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 95: - batch_foreach, 2, 3, 1>, 2, 3>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 96: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 97: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 98: - batch_foreach, 3, 3, 1>, 3, 3>(instr, locals, device); + batch_foreach, 2, 3, 1>, 2, 3>(instr, locals, device); break; case 99: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 100: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 101: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 3, 3, 1>, 3, 3>(instr, locals, device); break; case 102: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 103: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 104: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 105: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 106: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 107: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 108: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 109: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 110: - op_matrix_element(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 111: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 112: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 113: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 114: - op_matmul(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 115: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 116: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 117: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 118: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 119: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 120: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 121: - op_rqs_reshape(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 122: - batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 123: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 124: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 125: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); break; case 126: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 127: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 128: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 129: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 130: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 131: - op_discrete_histogram(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 132: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 133: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 134: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 135: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 136: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 137: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 138: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 139: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 140: - op_quantile(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 141: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 142: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 143: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + op_quantile(instr, locals, device); break; case 144: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 145: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 146: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 147: - op_nonzero(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 148: - op_batch_gather(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 149: - op_batch_scatter(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 150: - op_random(instr, locals, device); + op_nonzero(instr, locals, device); break; case 151: - op_random_int(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 152: - op_unweight(instr, locals, device); + op_batch_scatter(instr, locals, device); break; case 153: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_random(instr, locals, device); break; case 154: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_random_int(instr, locals, device); break; case 155: - op_vegas_histogram(instr, locals, device); + op_unweight(instr, locals, device); break; case 156: + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + break; +case 157: + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + break; +case 158: + op_vegas_histogram(instr, locals, device); + break; +case 159: op_histogram(instr, locals, device); break; diff --git a/madspace/src/kernels/multichannel.hpp b/madspace/src/kernels/multichannel.hpp index 836eef7f9..7b480849c 100644 --- a/madspace/src/kernels/multichannel.hpp +++ b/madspace/src/kernels/multichannel.hpp @@ -126,6 +126,66 @@ KERNELSPEC void kernel_apply_subchannel_weights( } } +template +KERNELSPEC void kernel_compress_channel_weights( + IIn channel_index, + FIn channel_weights, + IIn keep_count, + FOut chan_weight_values, + IOut chan_weight_indices +) { + IVal keep_index = channel_index; + std::size_t n_keep = chan_weight_values.size(); + for (std::size_t i = 0; i < n_keep - 1; i++) { + chan_weight_values[i] = 0.0; + chan_weight_indices[i] = -1; + } + chan_weight_values[n_keep - 1] = channel_weights.gather(keep_index); + chan_weight_indices[n_keep - 1] = keep_index; + if (n_keep == 1) { + return; + } + for (std::size_t i = 0; i < channel_weights.size(); i++) { + FVal value = where(i == keep_index, -1.0, channel_weights[i]); + BVal low_mask = value > chan_weight_values[0]; + chan_weight_values[0] = where(low_mask, value, chan_weight_values[0]); + chan_weight_indices[0] = where(low_mask, i, chan_weight_indices[0]); + + for (std::size_t j = 0; j < n_keep - 2; j++) { + FVal val1 = chan_weight_values[j], val2 = chan_weight_values[j + 1]; + IVal idx1 = chan_weight_indices[j], idx2 = chan_weight_indices[j + 1]; + BVal mask = val1 > val2; + chan_weight_values[j] = where(mask, val2, val1); + chan_weight_indices[j] = where(mask, idx2, idx1); + chan_weight_values[j + 1] = where(mask, val1, val2); + chan_weight_indices[j + 1] = where(mask, idx1, idx2); + } + } +} + +template +KERNELSPEC void kernel_restore_channel_weights( + FIn chan_weight_values, + IIn chan_weight_indices, + IIn full_count, + FOut channel_weights +) { + FVal epsilon = 1e-12; + for (std::size_t i = 0; i < channel_weights.size(); ++i) { + channel_weights[i] = epsilon; + } + FVal sum = channel_weights.size() * epsilon; + for (std::size_t i = 0; i < chan_weight_values.size(); ++i) { + FVal value = chan_weight_values[i]; + IVal index = chan_weight_indices[i]; + channel_weights.scatter_add(where(index == -1, 0, index), value); + sum += value; + } + for (std::size_t i = 0; i < channel_weights.size(); ++i) { + channel_weights[i] = channel_weights[i] / sum; + } +} + template KERNELSPEC void kernel_permute_momenta( FIn momenta, IIn permutations, IIn index, FOut output diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 2d808ee34..35a632bb0 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -29,7 +29,8 @@ Integrand::Integrand( const nested_vector2& active_flavors, const std::vector& flavor_remap, const std::vector& flavor_factors, - const std::vector& flavor_mirror + const std::vector& flavor_mirror, + std::size_t compressed_channel_weight_count ) : FunctionGenerator( "Integrand", @@ -48,14 +49,26 @@ Integrand::Integrand( ret_types.push_back("latent", batch_float_array(mapping.random_dim())); ret_types.push_back("adaptive_prob", batch_float); ret_types.push_back("channel_index", batch_int); - ret_types.push_back( - "channel_weights", - batch_float_array( - chan_weight_remap.size() > 0 ? remapped_chan_count - : subchan_weights ? subchan_weights->channel_count() - : diff_xs.matrix_element().diagram_count() - ) - ); + std::size_t madnis_channel_count = chan_weight_remap.size() > 0 + ? remapped_chan_count + : subchan_weights + ? subchan_weights->channel_count() + : diff_xs.matrix_element().diagram_count(); + if (madnis_channel_count > 1 && + madnis_channel_count * 8 > compressed_channel_weight_count * 12) { + ret_types.push_back( + "channel_weight_values", + batch_float_array(compressed_channel_weight_count) + ); + ret_types.push_back( + "channel_weight_indices", + batch_int_array(compressed_channel_weight_count) + ); + } else { + ret_types.push_back( + "channel_weights", batch_float_array(madnis_channel_count) + ); + } ret_types.push_back( "cwnet_input", batch_float_array( @@ -119,6 +132,7 @@ Integrand::Integrand( _chan_weight_net(chan_weight_net), _chan_weight_remap(chan_weight_remap), _remapped_chan_count(remapped_chan_count), + _compressed_channel_weight_count(compressed_channel_weight_count), _madnis_training(madnis_training), _drop_cuts_and_rescale(drop_cuts_and_rescale), _partial_weights(partial_weights), @@ -501,6 +515,10 @@ NamedVector Integrand::build_channel_part( Value adaptive_prob = adaptive_probs.empty() ? fb.full({1., batch_size_val}) : fb.product(adaptive_probs); + if (_drop_cuts_and_rescale) { + adaptive_prob = + fb.div(fb.accept_norm(indices_acc, adaptive_prob), adaptive_prob); + } NamedVector out; @@ -583,6 +601,8 @@ NamedVector Integrand::build_common_part( : _subchan_weights ? _subchan_weights->channel_count() : _diff_xs.matrix_element().diagram_count(); + bool use_compressed_channel_weights = + channel_count > 1 && channel_count * 8 > _compressed_channel_weight_count * 12; Value chan_weights_acc; if (channel_count > 1 && _prop_chan_weights) { chan_weights_acc = _prop_chan_weights->build_function(fb, {momenta_acc}).at(0); @@ -671,22 +691,46 @@ NamedVector Integrand::build_common_part( outputs.push_back("full_weight", optional_cut(full_weight)); outputs.push_back("weight", optional_cut(weight)); outputs.push_back("latent", optional_cut(args.at("latent"))); - Value norm = _drop_cuts_and_rescale - ? fb.accept_norm(indices_acc, args.at("adaptive_prob")) - : Value(1.); - outputs.push_back( - "adaptive_prob", fb.div(norm, optional_cut(args.at("adaptive_prob"))) - ); + outputs.push_back("adaptive_prob", optional_cut(args.at("adaptive_prob"))); outputs.push_back("channel_index", optional_cut(args.at("chan_index"))); if (channel_count > 1) { - auto cw_flat = fb.full( - {1. / channel_count, - batch_size_val, - static_cast(channel_count)} - ); - outputs.push_back( - "channel_weights", scatter_or_drop(cw_flat, prior_chan_weights_acc) - ); + if (use_compressed_channel_weights) { + Value chan_index_acc = + fb.batch_gather(indices_acc, args.at("chan_index")); + auto [chan_weight_values_acc, chan_weight_indices_acc] = + fb.compress_channel_weights( + chan_index_acc, + prior_chan_weights_acc, + static_cast(_compressed_channel_weight_count) + ); + auto cw_values_default = fb.full( + {0., + batch_size_val, + static_cast(_compressed_channel_weight_count)} + ); + auto cw_indices_default = fb.full( + {static_cast(-1), + batch_size_val, + static_cast(_compressed_channel_weight_count)} + ); + outputs.push_back( + "channel_weight_values", + scatter_or_drop(cw_values_default, chan_weight_values_acc) + ); + outputs.push_back( + "channel_weight_indices", + scatter_or_drop(cw_indices_default, chan_weight_indices_acc) + ); + } else { + auto cw_flat = fb.full( + {1. / channel_count, + batch_size_val, + static_cast(channel_count)} + ); + outputs.push_back( + "channel_weights", scatter_or_drop(cw_flat, prior_chan_weights_acc) + ); + } } else { outputs.push_back( "channel_weights", diff --git a/madspace/src/phasespace/madnis.cpp b/madspace/src/phasespace/madnis.cpp index e3411c4b8..3e2485e60 100644 --- a/madspace/src/phasespace/madnis.cpp +++ b/madspace/src/phasespace/madnis.cpp @@ -5,12 +5,16 @@ using namespace madspace; MadnisLoss::MadnisLoss( const std::vector>& functions, const std::optional& cwnet, - double softclip_threshold + double softclip_threshold, + std::size_t compressed_channel_weight_count ) : FunctionGenerator( "MadnisLoss", [&] { NamedVector arg_types; + bool use_compressed_channel_weights = cwnet && + cwnet->mlp().output_dim() > 1 && + cwnet->mlp().output_dim() * 8 > compressed_channel_weight_count * 12; for (std::size_t index = 0; auto& func : functions) { arg_types.push_back( std::format("chan{}_integrand", index), batch_float @@ -23,10 +27,21 @@ MadnisLoss::MadnisLoss( std::format("chan{}_cwnet_inputs", index), batch_float_array(cwnet->preprocessing().output_dim()) ); - arg_types.push_back( - std::format("chan{}_channel_weights", index), - batch_float_array(cwnet->mlp().output_dim()) - ); + if (use_compressed_channel_weights) { + arg_types.push_back( + std::format("chan{}_channel_weight_values", index), + batch_float_array(compressed_channel_weight_count) + ); + arg_types.push_back( + std::format("chan{}_channel_weight_indices", index), + batch_int_array(compressed_channel_weight_count) + ); + } else { + arg_types.push_back( + std::format("chan{}_channel_weights", index), + batch_float_array(cwnet->mlp().output_dim()) + ); + } arg_types.push_back( std::format("chan{}_channel_indices", index), batch_int ); @@ -45,22 +60,34 @@ MadnisLoss::MadnisLoss( ), _functions(functions), _cwnet(cwnet), - _softclip_threshold(softclip_threshold) {} + _softclip_threshold(softclip_threshold), + _compressed_channel_weight_count(compressed_channel_weight_count) {} NamedVector MadnisLoss::build_function_impl( FunctionBuilder& fb, const NamedVector& args ) const { ValueVec integrands, flow_probs, sample_probs, cwnet_inputs, cwnet_priors, chan_indices; - std::size_t extra_args = _cwnet ? 5 : 2; + bool use_compressed_channel_weights = _cwnet && _cwnet->mlp().output_dim() > 1 && + _cwnet->mlp().output_dim() * 8 > _compressed_channel_weight_count * 12; + std::size_t extra_args = _cwnet ? (use_compressed_channel_weights ? 6 : 5) : 2; for (std::size_t index = 0, arg_index = 0; auto& func : _functions) { std::size_t arg_index_end = arg_index + func->arg_types().size() + extra_args; integrands.push_back(args.at(arg_index)); sample_probs.push_back(args.at(arg_index + 1)); if (_cwnet) { cwnet_inputs.push_back(args.at(arg_index + 2)); - cwnet_priors.push_back(args.at(arg_index + 3)); - chan_indices.push_back(args.at(arg_index + 4)); + if (use_compressed_channel_weights) { + cwnet_priors.push_back(fb.restore_channel_weights( + args.at(arg_index + 3), + args.at(arg_index + 4), + static_cast(_cwnet->mlp().output_dim()) + )); + chan_indices.push_back(args.at(arg_index + 5)); + } else { + cwnet_priors.push_back(args.at(arg_index + 3)); + chan_indices.push_back(args.at(arg_index + 4)); + } } ValueVec func_args( args.begin() + arg_index + extra_args, args.begin() + arg_index_end diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp index 73fceb63d..611570449 100644 --- a/madspace/src/python/instruction_set.hpp +++ b/madspace/src/python/instruction_set.hpp @@ -34,6 +34,7 @@ void add_instructions(py::classh& 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")); @@ -96,6 +97,8 @@ void add_instructions(py::classh& fb) { fb.def("sde2_channel_weights", &FunctionBuilder::sde2_channel_weights, py::arg("invariants"), py::arg("masses"), py::arg("widths"), py::arg("indices")); fb.def("subchannel_weights", &FunctionBuilder::subchannel_weights, py::arg("invariants"), py::arg("masses"), py::arg("widths"), py::arg("indices"), py::arg("on_shell"), py::arg("group_sizes")); fb.def("apply_subchannel_weights", &FunctionBuilder::apply_subchannel_weights, py::arg("channel_weights_in"), py::arg("subchannel_weights"), py::arg("channel_indices"), py::arg("subchannel_indices")); + fb.def("compress_channel_weights", &FunctionBuilder::compress_channel_weights, py::arg("channel_index"), py::arg("channel_weights"), py::arg("keep_count")); + fb.def("restore_channel_weights", &FunctionBuilder::restore_channel_weights, py::arg("chan_weight_values"), py::arg("chan_weight_indices"), py::arg("full_count")); fb.def("pt_eta_phi_x", &FunctionBuilder::pt_eta_phi_x, py::arg("p_ext"), py::arg("x1"), py::arg("x2")); fb.def("mirror_momenta", &FunctionBuilder::mirror_momenta, py::arg("p_ext"), py::arg("mirror")); fb.def("momenta_to_x1x2", &FunctionBuilder::momenta_to_x1x2, py::arg("p_ext"), py::arg("e_cm")); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 72c029e7c..3e39b1bbb 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", @@ -1234,7 +1236,8 @@ PYBIND11_MODULE(_madspace_py, m) { const nested_vector2&, const std::vector&, const std::vector&, - const std::vector&>(), + const std::vector&, + std::size_t>(), py::arg("mapping"), py::arg("diff_xs"), py::arg("adaptive_map") = std::monostate{}, @@ -1255,7 +1258,8 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("active_flavors") = nested_vector2{}, py::arg("flavor_remap") = std::vector{}, py::arg("flavor_factors") = std::vector{}, - py::arg("flavor_mirror") = std::vector{} + py::arg("flavor_mirror") = std::vector{}, + py::arg("compressed_channel_weight_count") = 50 ) .def("particle_count", &Integrand::particle_count) .def("madnis_training", &Integrand::madnis_training) @@ -1288,10 +1292,12 @@ PYBIND11_MODULE(_madspace_py, m) { py::init< const std::vector>&, const std::optional&, - double>(), + double, + std::size_t>(), py::arg("functions"), py::arg("cwnet"), - py::arg("softclip_threshold") = 0.0 + py::arg("softclip_threshold") = 0.0, + py::arg("compressed_channel_weight_count") = 50 ); add_enum( @@ -1306,7 +1312,6 @@ PYBIND11_MODULE(_madspace_py, m) { py::classh(m, "MadnisConfig") .def(py::init<>()) - .def_readwrite("verbosity", &MadnisTraining::Config::verbosity) .def_readwrite("learning_rate", &MadnisTraining::Config::learning_rate) .def_readwrite("batches", &MadnisTraining::Config::batches) .def_readwrite("log_interval", &MadnisTraining::Config::log_interval) @@ -1349,6 +1354,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 @@ -1363,6 +1371,10 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def_readwrite( "softclip_threshold", &MadnisTraining::Config::softclip_threshold + ) + .def_readwrite( + "compressed_channel_weight_count", + &MadnisTraining::Config::compressed_channel_weight_count ); py::classh(m, "MadnisTraining") @@ -1383,19 +1395,37 @@ PYBIND11_MODULE(_madspace_py, m) { .def("active_channels", &MadnisTraining::active_channels) .def("active_channel_count", &MadnisTraining::active_channel_count); + py::classh(m, "StatusFile") + .def( + py::init(), + py::arg("file_name"), + py::arg("min_interval_sec") = 10.0 + ); + + py::classh(m, "TrainingArgs") + .def( + py::init< + const MadnisTraining::Config&, + const std::vector>&, + const std::optional&>(), + py::arg("config"), + py::arg("integrands"), + py::arg("cwnet") + ); + py::classh(m, "MultiMadnisTraining") .def( py::init< ContextPtr, ContextPtr, - const MadnisTraining::Config&, - const nested_vector2>&, - const std::vector>&>(), + const std::vector&, + Verbosity, + std::shared_ptr>(), py::arg("generator_context"), py::arg("optimizer_context"), - py::arg("config"), - py::arg("integrands"), - py::arg("cwnets") + py::arg("training_args"), + py::arg("verbosity"), + py::arg("status_file") = std::shared_ptr() ) .def("train", &MultiMadnisTraining::train) .def("active_channels", &MultiMadnisTraining::active_channels); @@ -1711,11 +1741,11 @@ PYBIND11_MODULE(_madspace_py, m) { py::init< const std::vector&, const std::vector>&, - const std::string&, + std::shared_ptr, const GeneratorConfig&>(), py::arg("contexts"), py::arg("channels"), - py::arg("status_file") = "", + py::arg("status_file") = std::shared_ptr(), py::arg_v( "config", EventGenerator::default_config, diff --git a/tests/unit_tests/various/test_banner.py b/tests/unit_tests/various/test_banner.py index 0afddca16..1b6e96842 100755 --- a/tests/unit_tests/various/test_banner.py +++ b/tests/unit_tests/various/test_banner.py @@ -1319,7 +1319,7 @@ def test_defaults_and_section_access(self): self.assertEqual(rc['generation']['events'], 100000) self.assertEqual(rc['beam']['e_cm'], 13000.0) self.assertIs(rc['vegas']['enable'], True) - self.assertIs(rc['madnis']['enable'], False) + self.assertEqual(rc['madnis']['enable'], 'auto') self.assertEqual(rc['phasespace']['mode'], 'multichannel') # the "enable" key colliding between [vegas] and [madnis] is stored # independently, so editing one does not affect the other