code-loader 1.0.189.dev0__tar.gz → 1.0.190.dev0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (37) hide show
  1. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/PKG-INFO +3 -3
  2. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/contract/datasetclasses.py +0 -20
  3. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/inner_leap_binder/leapbinder.py +1 -48
  4. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/inner_leap_binder/leapbinder_decorators.py +26 -241
  5. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/leaploader.py +1 -124
  6. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/leaploaderbase.py +0 -5
  7. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/pyproject.toml +1 -1
  8. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/LICENSE +0 -0
  9. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/README.md +0 -0
  10. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/__init__.py +0 -0
  11. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/contract/__init__.py +0 -0
  12. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/contract/enums.py +0 -0
  13. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/contract/exceptions.py +0 -0
  14. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/contract/mapping.py +0 -0
  15. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/contract/responsedataclasses.py +0 -0
  16. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/contract/sim_config.py +0 -0
  17. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/contract/visualizer_classes.py +0 -0
  18. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/default_losses.py +0 -0
  19. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/default_metrics.py +0 -0
  20. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/experiment_api/__init__.py +0 -0
  21. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/experiment_api/api.py +0 -0
  22. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/experiment_api/cli_config_utils.py +0 -0
  23. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/experiment_api/client.py +0 -0
  24. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/experiment_api/epoch.py +0 -0
  25. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/experiment_api/experiment.py +0 -0
  26. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/experiment_api/experiment_context.py +0 -0
  27. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/experiment_api/types.py +0 -0
  28. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/experiment_api/utils.py +0 -0
  29. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/experiment_api/workingspace_config_utils.py +0 -0
  30. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/inner_leap_binder/__init__.py +0 -0
  31. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/mixpanel_tracker.py +0 -0
  32. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/plot_functions/__init__.py +0 -0
  33. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/plot_functions/plot_functions.py +0 -0
  34. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/plot_functions/visualize.py +0 -0
  35. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/utils.py +0 -0
  36. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/visualizers/__init__.py +0 -0
  37. {code_loader-1.0.189.dev0 → code_loader-1.0.190.dev0}/code_loader/visualizers/default_visualizers.py +0 -0
@@ -1,8 +1,7 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.3
2
2
  Name: code-loader
3
- Version: 1.0.189.dev0
3
+ Version: 1.0.190.dev0
4
4
  Summary:
5
- Home-page: https://github.com/tensorleap/code-loader
6
5
  License: MIT
7
6
  Author: dorhar
8
7
  Author-email: doron.harnoy@tensorleap.ai
@@ -20,6 +19,7 @@ Requires-Dist: numpy (>=2.3.2,<3.0.0) ; python_version >= "3.11" and python_vers
20
19
  Requires-Dist: psutil (>=5.9.5,<6.0.0)
21
20
  Requires-Dist: pyyaml (>=6.0.2,<7.0.0)
22
21
  Requires-Dist: requests (>=2.32.3,<3.0.0)
22
+ Project-URL: Homepage, https://github.com/tensorleap/code-loader
23
23
  Project-URL: Repository, https://github.com/tensorleap/code-loader
24
24
  Description-Content-Type: text/markdown
25
25
 
@@ -92,16 +92,6 @@ SectionCallableInterface = Callable[[Union[int, str], PreprocessResponse], npt.N
92
92
  InstanceCallableInterface = Callable[[Union[int, str], PreprocessResponse, int], Optional[ElementInstance]]
93
93
  InstanceLengthCallableInterface = Callable[[Union[int, str], PreprocessResponse], int]
94
94
 
95
- # (sample_id, prev_inputs, prev_outputs, preprocess) -> next model inputs, or None to end the chain.
96
- # First call per chain receives prev_inputs=None, prev_outputs=None and returns the initial inputs.
97
- # Keys starting with '_' are passthrough state: never fed to the model, handed back verbatim in
98
- # prev_inputs on the next call. The hook must be stateless and deterministically seeded (the engine
99
- # may replay a step after crash recovery and relies on identical results).
100
- AutoregressiveStepCallableInterface = Callable[
101
- [Union[int, str], Optional[Dict[str, npt.NDArray[np.float32]]], Optional[Dict[str, npt.NDArray[np.float32]]],
102
- PreprocessResponse],
103
- Optional[Dict[str, npt.NDArray[np.float32]]]]
104
-
105
95
  MetadataSectionCallableInterface = Union[
106
96
  Callable[[Union[int, str], PreprocessResponse], int],
107
97
  Callable[[Union[int, str], PreprocessResponse], Dict[str, int]],
@@ -249,15 +239,6 @@ class CustomLatentSpaceHandler:
249
239
  function: SectionCallableInterface
250
240
  name: str = 'custom_latent_space'
251
241
 
252
-
253
- @dataclass
254
- class AutoregressiveStepHandler:
255
- function: AutoregressiveStepCallableInterface
256
- name: str = 'autoregressive_step'
257
- # Per-model-input shapes (unbatched, underscore passthrough keys excluded), discovered by running
258
- # the hook's first call at parse time. Fills the role InputHandler.shape plays for input encoders.
259
- input_shapes: Optional[Dict[str, List[int]]] = None
260
-
261
242
  @dataclass
262
243
  class PredictionTypeHandler:
263
244
  name: str
@@ -298,7 +279,6 @@ class DatasetIntegrationSetup:
298
279
  custom_layers: Dict[str, CustomLayerHandler] = field(default_factory=dict)
299
280
  custom_latent_space: Optional[CustomLatentSpaceHandler] = None
300
281
  simulations: List[SimulationHandler] = field(default_factory=list)
301
- autoregressive_step: Optional[AutoregressiveStepHandler] = None
302
282
 
303
283
 
304
284
  @dataclass
@@ -13,7 +13,7 @@ from code_loader.contract.datasetclasses import SectionCallableInterface, InputH
13
13
  CustomMultipleReturnCallableInterfaceMultiArgs, DatasetBaseHandler, custom_latent_space_attribute, \
14
14
  RawInputsForHeatmap, VisualizerHandlerData, MetricHandlerData, CustomLossHandlerData, SamplePreprocessResponse, \
15
15
  ElementInstanceMasksHandler, InstanceCallableInterface, CustomLatentSpaceHandler, InstanceMetricHandler, \
16
- SimulationHandler, _simulation_context, AutoregressiveStepHandler, AutoregressiveStepCallableInterface
16
+ SimulationHandler, _simulation_context
17
17
  from code_loader.contract.enums import LeapDataType, DataStateEnum, DataStateType, MetricDirection, DatasetMetadataType
18
18
  from code_loader.contract.mapping import NodeConnection, NodeMapping, NodeMappingType
19
19
  from code_loader.contract.responsedataclasses import DatasetTestResultPayload, LeapAnalysisConfiguration
@@ -524,26 +524,6 @@ class LeapBinder:
524
524
  """
525
525
  self.setup_container.custom_latent_space = CustomLatentSpaceHandler(function)
526
526
 
527
- def set_autoregressive_step(self, function: AutoregressiveStepCallableInterface) -> None:
528
- """
529
- Set the autoregressive step hook — the feedback function that drives a chain:
530
- it supplies the model's initial inputs on its first call (prev_inputs=None, prev_outputs=None)
531
- and turns each step's inputs/outputs into the next step's inputs until it returns None.
532
- An autoregressive integration has no input encoders; the hook is the sole input source.
533
- """
534
- if self.setup_container.autoregressive_step is not None:
535
- raise Exception('tensorleap_autoregressive_step is already defined. '
536
- 'Only one autoregressive step hook is allowed per integration.')
537
- self.setup_container.autoregressive_step = AutoregressiveStepHandler(function)
538
-
539
- # Builtin chain metadata, declared at parse time so it survives the reporter's
540
- # metadata type mapping; the placeholder values are overwritten by the engine when a
541
- # chain finalizes (realized length, truncated-by-safety-cap flag).
542
- def builtin_chain_metadata(idx: Any, preprocess: PreprocessResponse) -> Dict[str, Any]:
543
- return {'length': -1, 'truncated': False}
544
-
545
- self.set_metadata(builtin_chain_metadata, 'builtin_chain')
546
-
547
527
  def set_custom_layer(self, custom_layer: Type[Any], name: str, inspect_layer: bool = False,
548
528
  kernel_index: Optional[int] = None, use_custom_latent_space: bool = False) -> None:
549
529
  """
@@ -754,7 +734,6 @@ class LeapBinder:
754
734
  )
755
735
 
756
736
  def check(self) -> None:
757
- self.validate_autoregressive_setup()
758
737
  preprocess_result = self.get_preprocess_result()
759
738
  self.check_preprocess(preprocess_result)
760
739
  self.check_handlers(preprocess_result)
@@ -762,32 +741,6 @@ class LeapBinder:
762
741
  self.validate_ignore_latent_spaces()
763
742
  print("Successful!")
764
743
 
765
- def validate_autoregressive_setup(self) -> None:
766
- """Structural rules for autoregressive integrations (no-op when the hook is absent).
767
-
768
- Order-independent: runs after the whole script registered, because decorators may
769
- appear in any order in the integration file.
770
- """
771
- if self.setup_container.autoregressive_step is None:
772
- return
773
- if self.setup_container.inputs:
774
- raise Exception(
775
- "An autoregressive integration must not declare input encoders — "
776
- "tensorleap_autoregressive_step is the sole source of model inputs "
777
- "(its first call returns the initial inputs). Found input encoders: "
778
- f"{[handler.name for handler in self.setup_container.inputs]}.")
779
- if self.setup_container.instance_masks:
780
- raise Exception(
781
- "Element instances are not supported together with tensorleap_autoregressive_step: "
782
- "instance generation masks the sample's encoded inputs, which do not exist in an "
783
- "autoregressive integration. Remove the instance encoders or the autoregressive hook.")
784
- if self.setup_container.custom_latent_space is not None:
785
- raise Exception(
786
- "tensorleap_custom_latent_space is not supported together with "
787
- "tensorleap_autoregressive_step: the custom latent space function only sees "
788
- "(sample_id, preprocess) — it runs before generation and cannot observe the chain. "
789
- "The chain's latent is the final step's model latent.")
790
-
791
744
  def validate_ignore_latent_spaces(self) -> None:
792
745
  """Validate leap_analysis_configuration.ignore_latent_spaces against this binder.
793
746
 
@@ -21,7 +21,7 @@ from code_loader.contract.datasetclasses import CustomCallableInterfaceMultiArgs
21
21
  CustomMultipleReturnCallableInterfaceMultiArgs, ConfusionMatrixCallableInterfaceMultiArgs, CustomCallableInterface, \
22
22
  VisualizerCallableInterface, MetadataSectionCallableInterface, PreprocessResponse, SectionCallableInterface, \
23
23
  ConfusionMatrixElement, SamplePreprocessResponse, PredictionTypeHandler, InstanceCallableInterface, ElementInstance, \
24
- InstanceLengthCallableInterface, AutoregressiveStepCallableInterface
24
+ InstanceLengthCallableInterface
25
25
  from code_loader.contract.enums import MetricDirection, LeapDataType, DatasetMetadataType, DataStateType
26
26
  from code_loader import leap_binder, LeapLoader
27
27
  from code_loader.contract.mapping import NodeMapping, NodeMappingType, NodeConnection
@@ -197,11 +197,11 @@ def validate_output_structure(result, func_name: str, expected_type_name="np.nda
197
197
  def batch_warning(result, func_name):
198
198
  if len(result.shape) > 0 and result.shape[0] == 1:
199
199
  warnings.warn(
200
- f"{func_name} warning: Tensorleap will add a batch dimension at axis 0 to the output of {func_name}, "
201
- f"although the detected size of axis 0 is already 1. "
202
- f"This may lead to an extra batch dimension (e.g., shape (1, 1, ...)). "
203
- f"Please ensure that the output of '{func_name}' is not already batched "
204
- f"to avoid computation errors."
200
+ f"{func_name} warning: Tensorleap will add a batch dimension at axis 0 of "
201
+ f"{func_name}'s output, but axis 0 is already size 1 — this can produce a "
202
+ f"doubly-batched shape (e.g. (1, 1, ...)). If axis 0 is a real dimension "
203
+ f"(e.g. a single channel) and not an accidental batch, you can ignore this; "
204
+ f"only fix it if you pre-batched the output."
205
205
  )
206
206
 
207
207
 
@@ -349,6 +349,14 @@ def _safe_get_item(key):
349
349
 
350
350
 
351
351
  def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]] = _UNSET):
352
+ """Register the model-loading function.
353
+
354
+ ``prediction_types``: declare one ``PredictionTypeHandler`` per model output, in
355
+ output order. Multi-output models must declare an entry for EVERY output — the
356
+ count is validated against the model's outputs — so for an output you don't consume
357
+ as a prediction, pass a throwaway handler to satisfy the count. Omit the argument
358
+ only when the model has a single prediction output.
359
+ """
352
360
  prediction_types_was_provided = prediction_types is not _UNSET
353
361
 
354
362
  if not prediction_types_was_provided:
@@ -567,6 +575,9 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
567
575
  def get_inputs(self):
568
576
  return self.model.get_inputs()
569
577
 
578
+ def get_outputs(self):
579
+ return self.model.get_outputs()
580
+
570
581
  model_placeholder = ModelPlaceholder(prediction_types)
571
582
  if not _call_from_tl_platform:
572
583
  update_env_params_func("tensorleap_load_model", "v")
@@ -631,6 +642,13 @@ def tensorleap_load_model(prediction_types: Optional[List[PredictionTypeHandler]
631
642
 
632
643
  return FollowInputIndex()
633
644
 
645
+ def get_outputs(self):
646
+ # Mapping-mode counterpart of the ONNX-run get_outputs(): in mapping
647
+ # mode outputs are reached by index, and get_outputs()[i] resolves to
648
+ # Prediction{i} — exactly like run()/__call__, which also return a
649
+ # ModelOutputPlaceholder. Reusing it keeps the mapping unchanged.
650
+ return ModelOutputPlaceholder()
651
+
634
652
  return ModelPlaceholder()
635
653
 
636
654
  def final_inner(*args, **kwargs):
@@ -1628,232 +1646,6 @@ def tensorleap_custom_latent_space():
1628
1646
  return decorating_function
1629
1647
 
1630
1648
 
1631
- def tensorleap_autoregressive_step():
1632
- """The feedback hook that drives an autoregressive chain.
1633
-
1634
- Signature of the decorated function:
1635
- (sample_id, prev_inputs, prev_outputs, preprocess) -> Optional[Dict[str, np.ndarray]]
1636
-
1637
- First call per chain: prev_inputs=None, prev_outputs=None — returns the initial model inputs.
1638
- Later calls receive the previous step's model inputs and outputs (unbatched, per-sample) and
1639
- return the next step's full input dict; returning None ends the chain. Keys starting with '_'
1640
- are passthrough state: never fed to the model, handed back verbatim in prev_inputs next call.
1641
- The hook must be stateless and deterministically seeded (seed stochasticity from sample_id) —
1642
- the engine replays steps after crash recovery and relies on identical results.
1643
- """
1644
- def decorating_function(user_function: AutoregressiveStepCallableInterface):
1645
- # Contract: shapes and key set are fixed across all calls AND all chains (compiled graphs
1646
- # need static shapes), so one record shared across samples is the correct scope.
1647
- first_call_record: Dict[str, Any] = {'keys': None, 'shapes': None}
1648
-
1649
- def _normalize_args(args, kwargs):
1650
- expected_names = ["sample_id", "prev_inputs", "prev_outputs", "preprocess"]
1651
- validate_args_structure(
1652
- *args,
1653
- types_order=[Union[int, str], Union[dict, type(None)], Union[dict, type(None)],
1654
- PreprocessResponse],
1655
- func_name=user_function.__name__, expected_names=expected_names, **kwargs)
1656
- normalized = []
1657
- for i, name in enumerate(expected_names):
1658
- normalized.append(args[i] if i < len(args) else kwargs[name])
1659
- return normalized
1660
-
1661
- def _validate_input_args(sample_id, prev_inputs, prev_outputs, preprocess_response):
1662
- assert (prev_inputs is None) == (prev_outputs is None), \
1663
- (f'{user_function.__name__}() validation failed: '
1664
- f'prev_inputs and prev_outputs must both be None (first call) or both be dicts '
1665
- f'(subsequent calls). Got prev_inputs={type(prev_inputs).__name__}, '
1666
- f'prev_outputs={type(prev_outputs).__name__}.')
1667
- assert type(sample_id) == preprocess_response.sample_id_type, \
1668
- (f'{user_function.__name__}() validation failed: '
1669
- f'Argument sample_id should be as the same type as defined in the preprocess response '
1670
- f'{preprocess_response.sample_id_type}. Got {type(sample_id)}.')
1671
-
1672
- def _validate_result(result, is_first_call):
1673
- if result is None:
1674
- assert not is_first_call, \
1675
- (f'{user_function.__name__}() validation failed: '
1676
- f'the first call (prev_inputs=None, prev_outputs=None) must return the initial '
1677
- f'model inputs dict — returning None on the first call would produce an empty chain.')
1678
- return
1679
- assert isinstance(result, dict), \
1680
- (f'{user_function.__name__}() validation failed: '
1681
- f'the return value must be a dict of named model input tensors, or None to end the '
1682
- f'chain. Got {type(result)}.')
1683
- model_input_keys = [key for key in result if not key.startswith('_')]
1684
- assert model_input_keys, \
1685
- (f'{user_function.__name__}() validation failed: '
1686
- f'the returned dict contains no model inputs — every key starts with "_" '
1687
- f'(underscore keys are passthrough state, never fed to the model).')
1688
- for key, value in result.items():
1689
- assert isinstance(key, str), \
1690
- (f'{user_function.__name__}() validation failed: '
1691
- f'all keys in the returned dict must be strings. Got {type(key)}.')
1692
- if key.startswith('_'):
1693
- continue
1694
- assert isinstance(value, np.ndarray), \
1695
- (f'{user_function.__name__}() validation failed: '
1696
- f'value for model input "{key}" must be a numpy array. Got {type(value)}.')
1697
-
1698
- def _validate_consistency(result):
1699
- # Rules 2-3 of the autoregressive contract: identical key set and identical
1700
- # (non-passthrough) shapes on every call — compiled graphs need static shapes.
1701
- if result is None:
1702
- return
1703
- keys = set(result.keys())
1704
- shapes = {key: tuple(value.shape) for key, value in result.items()
1705
- if not key.startswith('_') and isinstance(value, np.ndarray)}
1706
- if first_call_record['keys'] is None:
1707
- first_call_record['keys'] = keys
1708
- first_call_record['shapes'] = shapes
1709
- return
1710
- assert keys == first_call_record['keys'], \
1711
- (f'{user_function.__name__}() validation failed: '
1712
- f'every call must return the same key set as the first call. '
1713
- f'First call keys: {sorted(first_call_record["keys"])}, this call: {sorted(keys)}.')
1714
- for key, shape in shapes.items():
1715
- assert shape == first_call_record['shapes'][key], \
1716
- (f'{user_function.__name__}() validation failed: '
1717
- f'tensor shapes must be identical on every call (fixed-length padding + mask for '
1718
- f'growing sequences). Input "{key}" was {first_call_record["shapes"][key]} on the '
1719
- f'first call and {shape} now.')
1720
-
1721
- def _assert_deterministic(first_result, second_result):
1722
- error_message = (
1723
- f'{user_function.__name__}() validation failed: '
1724
- f'the hook is not deterministic — two calls with identical arguments returned different '
1725
- f'outputs. Seed any stochasticity (e.g. initial noise) from sample_id: the engine '
1726
- f'replays steps after crash recovery and relies on identical results.')
1727
- if first_result is None or second_result is None:
1728
- assert first_result is None and second_result is None, error_message
1729
- return
1730
- assert set(first_result.keys()) == set(second_result.keys()), error_message
1731
- for key, value in first_result.items():
1732
- if key.startswith('_') or not isinstance(value, np.ndarray):
1733
- continue
1734
- assert np.array_equal(value, second_result[key]), error_message
1735
-
1736
- def _squeeze_test_batch_dim(tensors_dict):
1737
- # Inside the integration test the model works on batch-1 tensors (this wrapper adds the
1738
- # batch axis to returned inputs, and model outputs come back batched). The user's hook
1739
- # body must see the same unbatched per-sample tensors it will see at engine runtime, so
1740
- # strip the leading batch-1 axis before calling it.
1741
- if tensors_dict is None:
1742
- return None
1743
- squeezed = {}
1744
- for key, value in tensors_dict.items():
1745
- if not key.startswith('_') and isinstance(value, np.ndarray) \
1746
- and value.ndim > 0 and value.shape[0] == 1:
1747
- squeezed[key] = value[0]
1748
- else:
1749
- squeezed[key] = value
1750
- return squeezed
1751
-
1752
- def inner_without_validate(sample_id, prev_inputs, prev_outputs, preprocess_response):
1753
- global _called_from_inside_tl_decorator
1754
- _called_from_inside_tl_decorator += 1
1755
- try:
1756
- result = user_function(sample_id, prev_inputs, prev_outputs, preprocess_response)
1757
- finally:
1758
- _called_from_inside_tl_decorator -= 1
1759
- return result
1760
-
1761
- leap_binder.set_autoregressive_step(inner_without_validate)
1762
-
1763
- def inner(*args, **kwargs):
1764
- if not _call_from_tl_platform:
1765
- set_current('tensorleap_autoregressive_step')
1766
- sample_id, prev_inputs, prev_outputs, preprocess_response = _normalize_args(args, kwargs)
1767
-
1768
- is_top_level_in_test = (_called_from_inside_tl_decorator == 0
1769
- and _called_from_inside_tl_integration_test_decorator)
1770
- if is_top_level_in_test:
1771
- prev_inputs = _squeeze_test_batch_dim(prev_inputs)
1772
- prev_outputs = _squeeze_test_batch_dim(prev_outputs)
1773
-
1774
- _validate_input_args(sample_id, prev_inputs, prev_outputs, preprocess_response)
1775
-
1776
- result = inner_without_validate(sample_id, prev_inputs, prev_outputs, preprocess_response)
1777
-
1778
- if is_top_level_in_test:
1779
- # Rule 7: determinism double-call. An unseeded sampler fails here, on the user's
1780
- # machine, instead of silently forking a chain on engine crash-replay.
1781
- second_result = inner_without_validate(sample_id, prev_inputs, prev_outputs,
1782
- preprocess_response)
1783
- _assert_deterministic(result, second_result)
1784
-
1785
- _validate_result(result, is_first_call=prev_inputs is None)
1786
- _validate_consistency(result)
1787
-
1788
- if is_top_level_in_test and result is not None:
1789
- result = {key: (np.expand_dims(value, axis=0)
1790
- if not key.startswith('_') and isinstance(value, np.ndarray) else value)
1791
- for key, value in result.items()}
1792
-
1793
- if not _call_from_tl_platform:
1794
- update_env_params_func("tensorleap_autoregressive_step", "v")
1795
- return result
1796
-
1797
- class _MappingInputsPlaceholder:
1798
- """Mapping-mode stand-in for the hook's returned dict.
1799
-
1800
- Indexing it by a model-input name yields a placeholder whose node_mapping the model
1801
- placeholder tags Input0/Input1/... — exactly how input encoders wire the graph. One
1802
- NodeConnection is recorded per non-passthrough key, once.
1803
- """
1804
- def __init__(self):
1805
- self._items: Dict[str, Any] = {}
1806
-
1807
- def __getitem__(self, key):
1808
- assert isinstance(key, str), \
1809
- f'Expected a string model input name, got {type(key)} instead.'
1810
- if key not in self._items:
1811
- class TempMapping:
1812
- pass
1813
- ret = TempMapping()
1814
- ret.node_mapping = NodeMapping(key, NodeMappingType.Input)
1815
- if not key.startswith('_'):
1816
- leap_binder.mapping_connections.append(NodeConnection(ret.node_mapping, None))
1817
- self._items[key] = ret
1818
- return self._items[key]
1819
-
1820
- def get(self, key, default=None):
1821
- return self[key]
1822
-
1823
- def _unsupported_iteration(self):
1824
- raise Exception(
1825
- 'Iterating over the autoregressive hook result is not supported inside the '
1826
- 'integration test. Pass the model its inputs with explicit keys, e.g. '
1827
- "model.run(None, {'x': inputs['x'], 't': inputs['t']}).")
1828
-
1829
- def items(self):
1830
- self._unsupported_iteration()
1831
-
1832
- def keys(self):
1833
- self._unsupported_iteration()
1834
-
1835
- def values(self):
1836
- self._unsupported_iteration()
1837
-
1838
- mapping_placeholder = _MappingInputsPlaceholder()
1839
-
1840
- def mapping_inner(*args, **kwargs):
1841
- return mapping_placeholder
1842
-
1843
- def final_inner(*args, **kwargs):
1844
- if os.environ.get(mapping_runtime_mode_env_var_mame):
1845
- return mapping_inner(*args, **kwargs)
1846
- else:
1847
- return inner(*args, **kwargs)
1848
-
1849
- if not _call_from_tl_platform:
1850
- register_decorator_row("tensorleap_autoregressive_step")
1851
-
1852
- return final_inner
1853
-
1854
- return decorating_function
1855
-
1856
-
1857
1649
  def tensorleap_preprocess():
1858
1650
  def decorating_function(user_function: Callable[[], List[PreprocessResponse]]):
1859
1651
  leap_binder.set_preprocess(user_function)
@@ -2684,12 +2476,6 @@ def tensorleap_status_table():
2684
2476
  table.append({"name": "tensorleap simulation (optional)", "Added to integration": UNKNOWN})
2685
2477
  _sim_tracking[sim_name] = "registered"
2686
2478
 
2687
- def register_decorator_row(name: str):
2688
- # Rows for decorators that only apply to some integration kinds (e.g. the autoregressive
2689
- # step hook) are appended on registration, so unrelated integrations never show them.
2690
- if _find_row(_remove_suffix(name, " (optional)")) is None:
2691
- table.append({"name": name, "Added to integration": UNKNOWN})
2692
-
2693
2479
  def mark_sim_result(sim_name: str, passed: bool):
2694
2480
  _sim_tracking[sim_name] = "passed" if passed else "failed"
2695
2481
 
@@ -2736,12 +2522,11 @@ def tensorleap_status_table():
2736
2522
  atexit.register(run_on_exit)
2737
2523
  sys.excepthook = handle_exception
2738
2524
 
2739
- return set_current, update_env_params, register_sim, mark_sim_result, register_decorator_row
2525
+ return set_current, update_env_params, register_sim, mark_sim_result
2740
2526
 
2741
2527
 
2742
2528
  if not _call_from_tl_platform:
2743
- set_current, update_env_params_func, register_sim, mark_sim_result, register_decorator_row = \
2744
- tensorleap_status_table()
2529
+ set_current, update_env_params_func, register_sim, mark_sim_result = tensorleap_status_table()
2745
2530
 
2746
2531
 
2747
2532
 
@@ -320,16 +320,12 @@ class LeapLoader(LeapLoaderBase):
320
320
  try:
321
321
  self.exec_script()
322
322
 
323
- global_leap_binder.validate_autoregressive_setup()
324
323
  preprocess_test_payload = self._check_preprocess()
325
324
  test_payloads.append(preprocess_test_payload)
326
325
  handlers_test_payloads = self._check_handlers()
327
326
  test_payloads.extend(handlers_test_payloads)
328
327
  simulation_test_payloads = self._check_simulations()
329
328
  test_payloads.extend(simulation_test_payloads)
330
- autoregressive_test_payload = self._check_autoregressive_step()
331
- if autoregressive_test_payload is not None:
332
- test_payloads.append(autoregressive_test_payload)
333
329
  is_valid = all([payload.is_passed for payload in test_payloads])
334
330
  setup_response = self.get_dataset_setup_response(handlers_test_payloads)
335
331
 
@@ -449,69 +445,6 @@ class LeapLoader(LeapLoaderBase):
449
445
  result_payloads.append(test_result)
450
446
  return result_payloads
451
447
 
452
- def _check_autoregressive_step(self) -> Optional[DatasetTestResultPayload]:
453
- handler = global_leap_binder.setup_container.autoregressive_step
454
- if handler is None:
455
- return None
456
- test_result = DatasetTestResultPayload(handler.name)
457
- try:
458
- preprocess_result = self._preprocess_result()
459
- input_shapes: Dict[str, List[int]] = {}
460
- for state, preprocess_response in preprocess_result.items():
461
- if preprocess_response.sample_ids_to_instance_mappings:
462
- raise Exception('Element instances are not supported together with '
463
- 'tensorleap_autoregressive_step.')
464
- sample_id = preprocess_response.sample_ids[0]
465
- first_result = handler.function(sample_id, None, None, preprocess_response)
466
- second_result = handler.function(sample_id, None, None, preprocess_response)
467
- if not isinstance(first_result, dict):
468
- raise Exception('The autoregressive step hook must return a dict of initial model '
469
- f'inputs on its first call, got {type(first_result)} '
470
- f'(state: {state.name}).')
471
- model_input_items = {key: value for key, value in first_result.items()
472
- if not key.startswith('_')}
473
- if not model_input_items:
474
- raise Exception('The autoregressive step hook returned no model inputs on its first '
475
- 'call — every key starts with "_" (underscore keys are passthrough '
476
- f'state, never fed to the model). State: {state.name}.')
477
- for key, value in model_input_items.items():
478
- if not isinstance(value, np.ndarray):
479
- raise Exception(f'The autoregressive step hook returned a non-numpy value for '
480
- f'model input "{key}" ({type(value)}). State: {state.name}.')
481
- if key in input_shapes and input_shapes[key] != list(value.shape):
482
- raise Exception(f'The autoregressive step hook returned different shapes for '
483
- f'model input "{key}" across states: {input_shapes[key]} vs '
484
- f'{list(value.shape)}. Shapes must be fixed.')
485
- input_shapes[key] = list(value.shape)
486
- if not np.array_equal(value, second_result[key]):
487
- raise Exception('The autoregressive step hook is not deterministic — two calls '
488
- 'with identical arguments returned different outputs for model '
489
- f'input "{key}". Seed any stochasticity from sample_id: the '
490
- 'engine replays steps after crash recovery and relies on '
491
- 'identical results.')
492
- handler.input_shapes = input_shapes
493
- except Exception as e:
494
- line_number, file_name, stacktrace = get_root_exception_file_and_line_number()
495
- test_result.display[TestingSectionEnum.Errors.name] = \
496
- f"{repr(e)} in file {file_name}, line_number: {line_number}\nStacktrace:\n{stacktrace}"
497
- test_result.is_passed = False
498
- return test_result
499
-
500
- def _ensure_autoregressive_input_shapes(self) -> Dict[str, List[int]]:
501
- handler = global_leap_binder.setup_container.autoregressive_step
502
- assert handler is not None
503
- if handler.input_shapes is None:
504
- preprocess_result = self._preprocess_result()
505
- first_state_response = next(iter(preprocess_result.values()))
506
- first_result = handler.function(first_state_response.sample_ids[0], None, None,
507
- first_state_response)
508
- if not isinstance(first_result, dict):
509
- raise Exception('The autoregressive step hook must return a dict of initial model '
510
- f'inputs on its first call, got {type(first_result)}.')
511
- handler.input_shapes = {key: list(value.shape) for key, value in first_result.items()
512
- if not key.startswith('_') and isinstance(value, np.ndarray)}
513
- return handler.input_shapes
514
-
515
448
  def run_simulation(self, sim_name, params=None, n_samples=1, seed=0,
516
449
  sample_ids=None, extend_preprocess=True):
517
450
  # type: (str, Optional[Dict[str, Any]], int, int, Optional[List[str]], bool) -> Dict[str, Any]
@@ -702,12 +635,6 @@ class LeapLoader(LeapLoaderBase):
702
635
  raise Exception(f"cant calculate shape for input, input name:{inp.name}")
703
636
  inputs.append(DatasetInputInstance(name=inp.name, shape=inp.shape, channel_dim=inp.channel_dim))
704
637
 
705
- if setup.autoregressive_step is not None:
706
- # The hook is the sole input source; report its step-0 inputs as the dataset inputs so
707
- # graph building, input-coverage validation and type dictionaries see real specs.
708
- for input_name, input_shape in self._ensure_autoregressive_input_shapes().items():
709
- inputs.append(DatasetInputInstance(name=input_name, shape=input_shape, channel_dim=-1))
710
-
711
638
  ground_truths = []
712
639
  for gt in setup.ground_truths:
713
640
  if gt.shape is None:
@@ -815,21 +742,7 @@ class LeapLoader(LeapLoaderBase):
815
742
  return result_agg
816
743
 
817
744
  def _get_inputs(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
818
- inputs = self._get_dataset_handlers(global_leap_binder.setup_container.inputs, state, sample_id)
819
- autoregressive_handler = global_leap_binder.setup_container.autoregressive_step
820
- if autoregressive_handler is not None:
821
- # Autoregressive integrations have no input encoders — the hook's first call supplies
822
- # the sample's (step-0) model inputs, so every sample-fetching flow (graph validation,
823
- # analysis, import checks) sees a fully-populated inputs dict.
824
- preprocess_response = self._preprocess_result()[state]
825
- step_zero_inputs = autoregressive_handler.function(sample_id, None, None, preprocess_response)
826
- if not isinstance(step_zero_inputs, dict):
827
- raise Exception(
828
- 'The autoregressive step hook must return a dict of initial model inputs on its '
829
- f'first call, got {type(step_zero_inputs)} for sample {sample_id}.')
830
- inputs.update({key: value for key, value in step_zero_inputs.items()
831
- if not key.startswith('_')})
832
- return inputs
745
+ return self._get_dataset_handlers(global_leap_binder.setup_container.inputs, state, sample_id)
833
746
 
834
747
  def _get_instances_masks(self, state: DataStateEnum, sample_id: Union[int, str], instance_id: int) -> Optional[Dict[str, ElementInstance]]:
835
748
  if instance_id is None:
@@ -852,12 +765,6 @@ class LeapLoader(LeapLoaderBase):
852
765
  def _get_gt(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
853
766
  return self._get_dataset_handlers(global_leap_binder.setup_container.ground_truths, state, sample_id)
854
767
 
855
- def get_gt(self, state: DataStateEnum, sample_id: Union[int, str]) -> Dict[str, npt.NDArray[np.float32]]:
856
- """Ground-truth tensors only ({} when no GT encoders are declared) — lets the engine's
857
- rollout fetch a chain's GT without re-running the autoregressive hook via get_sample."""
858
- self.exec_script()
859
- return self._get_gt(state, sample_id)
860
-
861
768
  @lru_cache()
862
769
  def _metadata_name_to_type(self) -> Dict[str, DatasetMetadataType]:
863
770
  global_leap_binder.check_preprocess(self._preprocess_result())
@@ -953,36 +860,6 @@ class LeapLoader(LeapLoaderBase):
953
860
  self.exec_script()
954
861
  return global_leap_binder.setup_container.custom_latent_space is not None
955
862
 
956
- @lru_cache()
957
- def has_autoregressive_step(self) -> bool:
958
- self.exec_script()
959
- return global_leap_binder.setup_container.autoregressive_step is not None
960
-
961
- def run_autoregressive_step(self, sample_id: Union[int, str],
962
- prev_inputs: Optional[Dict[str, npt.NDArray[np.float32]]],
963
- prev_outputs: Optional[Dict[str, npt.NDArray[np.float32]]],
964
- state: DataStateEnum) -> Optional[Dict[str, npt.NDArray[np.float32]]]:
965
- """Execute the user's autoregressive step hook against the in-process preprocess result.
966
-
967
- This is the engine rollout's per-step entry point (invoked on the generic pod). The first
968
- call per chain passes prev_inputs=None, prev_outputs=None and returns the initial model
969
- inputs; a None return ends the chain.
970
- """
971
- self.exec_script()
972
- handler = global_leap_binder.setup_container.autoregressive_step
973
- if handler is None:
974
- raise Exception('run_autoregressive_step was called but the integration has no '
975
- 'tensorleap_autoregressive_step hook.')
976
- preprocess_response = self._preprocess_result()[state]
977
- return handler.function(sample_id, prev_inputs, prev_outputs, preprocess_response)
978
-
979
- def get_prediction_names_in_order(self) -> List[str]:
980
- """Model output names by declared prediction-type order (validated against the model's
981
- output count at integration-test time), used to key prev_outputs for the hook."""
982
- self.exec_script()
983
- return [prediction_type.name for prediction_type in
984
- global_leap_binder.setup_container.prediction_types]
985
-
986
863
  def get_instances_data(self, state: DataStateEnum) -> Tuple[Dict[str, List[str]], Dict[str, str]]:
987
864
  preprocess_result = self._preprocess_result()
988
865
  preprocess_state = preprocess_result.get(state)
@@ -149,11 +149,6 @@ class LeapLoaderBase:
149
149
  def has_custom_latent_space_decorator(self) -> bool:
150
150
  pass
151
151
 
152
- def has_autoregressive_step(self) -> bool:
153
- # Default False so existing LeapLoaderBase implementations stay valid; the concrete
154
- # loaders that can answer (LeapLoader in-process, LeapLoaderWithRedis via RPC) override.
155
- return False
156
-
157
152
  @abstractmethod
158
153
  def get_heatmap_visualizer_raw_vis_input_arg_name(self, visualizer_name: str) -> Optional[str]:
159
154
  pass
@@ -1,7 +1,7 @@
1
1
  [tool.poetry]
2
2
  name = "code-loader"
3
3
 
4
- version = "1.0.189.dev0"
4
+ version = "1.0.190.dev0"
5
5
  description = ""
6
6
  authors = ["dorhar <doron.harnoy@tensorleap.ai>"]
7
7
  license = "MIT"