BatchRewardFunctionWrapper reward

@synalinks_export("synalinks.rewards.BatchRewardFunctionWrapper")
class BatchRewardFunctionWrapper(BatchReward):
    """Wrap a stateless batched function into a ``BatchReward``.

    The wrapped function receives the full batch and must return a
    ``list[float]`` of length ``batch_size``. It must be declared with
    ``async def``, since reward functions are awaited.

    Unlike per-sample functions, a batched function is never auto-wrapped by
    ``compile``: its ``batch -> list[float]`` signature cannot be told apart
    from a per-sample one, so it always has to be passed wrapped in this class.

    Example:

    ```python
    async def my_batch_reward(y_true, y_pred):
        # y_true, y_pred: list[JsonDataModel] of length batch_size
        return [1.0 if t.get_json() == p.get_json() else 0.0
                for t, p in zip(y_true, y_pred)]

    program.compile(
        reward=synalinks.rewards.BatchRewardFunctionWrapper(fn=my_batch_reward),
        optimizer=synalinks.optimizers.RandomFewShot(),
    )
    ```

    Args:
        fn (callable): Async batched reward function with signature
            ``fn(y_true, y_pred, **kwargs) -> list[float]``.
        name (str): Optional. string name of the reward instance.
        reduction (str): Optional. One of ``"mean"``, ``"sum"``, ``"min"``,
            ``"max"``, ``"none"`` or ``None``. Used by standalone
            ``__call__`` and propagated through ``compile`` to set the
            scalar reduction used by the trainer's progress log and the
            optimizer's candidate scoring (``"none"``/``None`` falls back
            to ``"mean"`` there).
        in_mask (list): Optional.
        out_mask (list): Optional.
        in_mask_pattern (str): Optional.
        out_mask_pattern (str): Optional.
        **kwargs (keyword arguments): Extra keyword arguments forwarded
            to ``fn``.
    """

    def __init__(
        self,
        fn,
        reduction="mean",
        name=None,
        in_mask=None,
        out_mask=None,
        in_mask_pattern=None,
        out_mask_pattern=None,
        **kwargs,
    ):
        super().__init__(
            name=name,
            reduction=reduction,
            in_mask=in_mask,
            out_mask=out_mask,
            in_mask_pattern=in_mask_pattern,
            out_mask_pattern=out_mask_pattern,
        )
        check_async_reward_fn(fn, self.__class__.__name__)
        self.fn = fn
        self._fn_kwargs = kwargs

    async def call(self, y_true, y_pred):
        return await self.fn(y_true, y_pred, **self._fn_kwargs)

    def get_config(self):
        config = super().get_config()
        config["fn"] = serialization_lib.serialize_synalinks_object(self.fn)
        config["fn_kwargs"] = serialization_lib.serialize_synalinks_object(
            self._fn_kwargs
        )
        return config

    @classmethod
    def from_config(cls, config):
        if "fn" in config:
            config = serialization_lib.deserialize_synalinks_object(config)
        fn_kwargs = config.pop("fn_kwargs", None) or {}
        return cls(**config, **fn_kwargs)

    def __repr__(self):
        return f"<BatchRewardFunctionWrapper({self.fn}, kwargs={self._fn_kwargs})>"