RewardFunctionWrapper reward

@synalinks_export("synalinks.rewards.RewardFunctionWrapper")
class RewardFunctionWrapper(Reward):
    """Wrap a stateless function into a `Reward`.

    You can use this to quickly build a reward from a function. The function needs
    to have the signature `fn(y_true, y_pred)` and to be declared with
    `async def`, since reward functions are awaited: a synchronous one raises a
    `TypeError` here rather than failing later inside the training loop.

    Example:

    ```python
    async def my_reward(y_true, y_pred):
        # ...
        return reward

    program.compile(
        reward=synalinks.rewards.RewardFunctionWrapper(fn=my_reward),
        optimizer=synalinks.optimizers.RandomFewShot(),
    )
    ```

    Wrapping is optional: `compile(reward=my_reward)` accepts the bare function
    and wraps it for you, naming the reward after the function. Reach for this
    class explicitly when you need masks, a custom `reduction`, or extra keyword
    arguments forwarded to `fn`:

    ```python
    async def length_under(y_true, y_pred, limit=100):
        return 1.0 if len(y_pred.get("answer")) < limit else 0.0

    program.compile(
        reward=synalinks.rewards.RewardFunctionWrapper(
            fn=length_under,
            limit=200,
            in_mask=["answer"],
        ),
        optimizer=synalinks.optimizers.RandomFewShot(),
    )
    ```

    Args:
        fn (callable): Async reward function to wrap, with signature
            ``fn(y_true, y_pred, **kwargs)``.
        name (str): Optional. string name of the reward instance.
        in_mask (list): Optional. list of keys to keep to compute the reward.
        out_mask (list): Optional. list of keys to remove to compute the reward.
        in_mask_pattern (str): Optional. Regex pattern; fields whose names match
            are kept (combined with ``in_mask`` via OR).
        out_mask_pattern (str): Optional. Regex pattern; fields whose names match
            are dropped (combined with ``out_mask`` via OR).
        **kwargs (keyword arguments): Keyword arguments to pass on 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)
        # Keep fn kwargs under their own key so they cannot collide with
        # base-class fields like ``name`` or ``reduction``.
        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"<RewardFunctionWrapper({self.fn}, kwargs={self._fn_kwargs})>"