zea 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

zea/__init__.py

@@ -7,7 +7,7 @@ from . import log
 
 # dynamically add __version__ attribute (see pyproject.toml)
 # __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.5"
+__version__ = "0.0.6"
 
 
 def _bootstrap_backend():

zea/agent/selection.py

@@ -11,11 +11,14 @@ For a comprehensive example usage, see: :doc:`../notebooks/agent/agent_example`
 All strategies are stateless, meaning that they do not maintain any internal state.
 """
 
+from typing import Callable
+
 import keras
 from keras import ops
 
 from zea import tensor_ops
 from zea.agent import masks
+from zea.backend.autograd import AutoGrad
 from zea.internal.registry import action_selection_registry
 
 
@@ -493,3 +496,166 @@ class CovarianceSamplingLines(LinesActionModel):
         best_mask = ops.squeeze(best_mask, axis=0)
 
         return best_mask, self.lines_to_im_size(best_mask)
+
+
+class TaskBasedLines(GreedyEntropy):
+    """Task-based line selection for maximizing information gain.
+
+    This action selection strategy chooses lines to maximize information gain with respect
+    to a downstream task outcome. It uses gradient-based saliency to identify which image
+    regions contribute most to task uncertainty, then selects lines accordingly.
+    """
+
+    def __init__(
+        self,
+        n_actions: int,
+        n_possible_actions: int,
+        img_width: int,
+        img_height: int,
+        downstream_task_function: Callable,
+        mean: float = 0,
+        std_dev: float = 1,
+        num_lines_to_update: int = 5,
+        **kwargs,
+    ):
+        """Initialize the TaskBasedLines action selection model.
+
+        Args:
+            n_actions (int): The number of actions the agent can take.
+            n_possible_actions (int): The number of possible actions (line positions).
+            img_width (int): The width of the input image.
+            img_height (int): The height of the input image.
+            downstream_task_function (Callable): A differentiable function that takes a
+                batch of inputs and produces scalar outputs. This represents the downstream
+                task for which information gain should be maximized.
+            mean (float, optional): The mean of the RBF used for reweighting. Defaults to 0.
+            std_dev (float, optional): The standard deviation of the RBF used for reweighting.
+                Defaults to 1.
+            num_lines_to_update (int, optional): The number of lines around the selected line
+                to update during reweighting. Must be odd. Defaults to 5.
+            **kwargs: Additional keyword arguments passed to the parent class.
+        """
+        super().__init__(
+            n_actions,
+            n_possible_actions,
+            img_width,
+            img_height,
+            mean,
+            std_dev,
+            num_lines_to_update,
+        )
+        self.downstream_task_function = downstream_task_function
+
+    def compute_output_and_saliency_propagation(self, particles):
+        """Compute saliency-weighted posterior variance for task-based selection.
+
+        This method computes how much each pixel contributes to the variance of the
+        downstream task output. It uses automatic differentiation to compute gradients
+        of the task function with respect to each particle, then weights the posterior
+        variance by the squared mean gradient.
+
+        Args:
+            particles (Tensor): Particles of shape (batch_size, n_particles, height, width)
+                representing the posterior distribution over images.
+
+        Returns:
+            Tensor: Pixelwise contribution to downstream task variance,
+                of shape (batch_size, height, width). Higher values indicate pixels
+                that contribute more to task uncertainty.
+        """
+        autograd = AutoGrad()
+
+        autograd.set_function(self.downstream_task_function)
+        downstream_grad_and_value_fn = autograd.get_gradient_and_value_jit_fn()
+        jacobian, _ = ops.vectorized_map(
+            lambda p: ops.vectorized_map(
+                downstream_grad_and_value_fn,
+                p,
+            ),
+            particles,
+        )
+
+        posterior_variance = ops.var(particles, axis=1)
+        mean_jacobian = ops.mean(jacobian, axis=1)
+        return posterior_variance * (mean_jacobian**2)
+
+    def sum_neighbouring_columns_into_n_possible_actions(self, full_linewise_salience):
+        """Aggregate column-wise saliency into line-wise saliency scores.
+
+        This method groups neighboring columns together to create saliency scores
+        for each possible line action. Since each line action may correspond to
+        multiple image columns, this aggregation is necessary to match the action space.
+
+        Args:
+            full_linewise_salience (Tensor): Saliency values for each column,
+                of shape (batch_size, full_image_width).
+
+        Returns:
+            Tensor: Aggregated saliency scores for each possible action,
+                of shape (batch_size, n_possible_actions).
+
+        Raises:
+            AssertionError: If the image width is not evenly divisible by n_possible_actions.
+        """
+        batch_size = ops.shape(full_linewise_salience)[0]
+        full_image_width = ops.shape(full_linewise_salience)[1]
+        assert full_image_width % self.n_possible_actions == 0, (
+            "n_possible_actions must divide evenly into image width"
+        )
+        cols_per_action = full_image_width // self.n_possible_actions
+        stacked_linewise_salience = ops.reshape(
+            full_linewise_salience,
+            (batch_size, self.n_possible_actions, cols_per_action),
+        )
+        return ops.sum(stacked_linewise_salience, axis=2)
+
+    def sample(self, particles):
+        """Sample actions using task-based information gain maximization.
+
+        This method computes which lines would provide the most information about
+        the downstream task by:
+        1. Computing pixelwise contribution to task variance using gradients
+        2. Aggregating contributions into line-wise scores
+        3. Greedily selecting lines with highest contribution scores
+        4. Reweighting scores around selected lines (inherited from GreedyEntropy)
+
+        Args:
+            particles (Tensor): Particles representing the posterior distribution,
+                of shape (batch_size, n_particles, height, width).
+
+        Returns:
+            Tuple[Tensor, Tensor, Tensor]:
+                - selected_lines_k_hot: Selected lines as k-hot vectors,
+                  shaped (batch_size, n_possible_actions)
+                - masks: Binary masks of shape (batch_size, img_height, img_width)
+                - pixelwise_contribution_to_var_dst: Pixelwise contribution to downstream
+                  task variance, of shape (batch_size, height, width)
+
+        Note:
+            Unlike the parent GreedyEntropy class, this method returns an additional
+            tensor containing the pixelwise contribution scores for analysis.
+        """
+        pixelwise_contribution_to_var_dst = self.compute_output_and_saliency_propagation(particles)
+        linewise_contribution_to_var_dst = ops.sum(pixelwise_contribution_to_var_dst, axis=1)
+        actionwise_contribution_to_var_dst = self.sum_neighbouring_columns_into_n_possible_actions(
+            linewise_contribution_to_var_dst
+        )
+
+        # Greedily select best line, reweight entropies, and repeat
+        all_selected_lines = []
+        for _ in range(self.n_actions):
+            max_contribution_line, actionwise_contribution_to_var_dst = ops.vectorized_map(
+                self.select_line_and_reweight_entropy,
+                actionwise_contribution_to_var_dst,
+            )
+            all_selected_lines.append(max_contribution_line)
+
+        selected_lines_k_hot = ops.any(
+            ops.one_hot(all_selected_lines, self.n_possible_actions, dtype=masks._DEFAULT_DTYPE),
+            axis=0,
+        )
+        return (
+            selected_lines_k_hot,
+            self.lines_to_im_size(selected_lines_k_hot),
+            pixelwise_contribution_to_var_dst,
+        )

zea/backend/__init__.py

@@ -25,6 +25,8 @@ Key Features
 
 """
 
+from contextlib import nullcontext
+
 import keras
 
 from zea import log
@@ -114,3 +116,90 @@ def _jit_compile(func, jax=True, tensorflow=True, **kwargs):
         log.warning("Initialize zea.Pipeline with jit_options=None to suppress this warning.")
         log.warning("Falling back to non-compiled mode.")
         return func
+
+
+class on_device:
+    """Context manager to set the device regardless of backend.
+
+    For the `torch` backend, you need to manually move the model and data to the device before
+    using this context manager.
+
+    Args:
+        device (str): Device string, e.g. ``'cuda'``, ``'gpu'``, or ``'cpu'``.
+
+    Example:
+        .. code-block:: python
+
+            with zea.backend.on_device("gpu:3"):
+                pipeline = zea.Pipeline([zea.keras_ops.Abs()])
+                output = pipeline(data=keras.random.normal((10, 10)))  # output is on "cuda:3"
+    """
+
+    def __init__(self, device: str):
+        self.device = self.get_device(device)
+        self._context = self.get_context(self.device)
+
+    def get_context(self, device):
+        if device is None:
+            return nullcontext()
+
+        if keras.backend.backend() == "tensorflow":
+            import tensorflow as tf
+
+            return tf.device(device)
+
+        if keras.backend.backend() == "jax":
+            import jax
+
+            return jax.default_device(device)
+        if keras.backend.backend() == "torch":
+            import torch
+
+            return torch.device(device)
+
+        return nullcontext()
+
+    def get_device(self, device: str):
+        if device is None:
+            return None
+
+        device = device.lower()
+
+        if keras.backend.backend() == "tensorflow":
+            return device.replace("cuda", "gpu")
+
+        if keras.backend.backend() == "jax":
+            from zea.backend.jax import str_to_jax_device
+
+            device = device.replace("cuda", "gpu")
+            return str_to_jax_device(device)
+
+        if keras.backend.backend() == "torch":
+            return device.replace("gpu", "cuda")
+
+    def __enter__(self):
+        self._context.__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._context.__exit__(exc_type, exc_val, exc_tb)
+
+
+if keras.backend.backend() in ["tensorflow", "jax", "numpy"]:
+
+    def func_on_device(func, device, *args, **kwargs):
+        """Moves all tensor arguments of a function to a specified device before calling it.
+
+        Args:
+            func (callable): Function to be called.
+            device (str): Device to move tensors to.
+            *args: Positional arguments to be passed to the function.
+            **kwargs: Keyword arguments to be passed to the function.
+        Returns:
+            The output of the function.
+        """
+        with on_device(device):
+            return func(*args, **kwargs)
+elif keras.backend.backend() == "torch":
+    from zea.backend.torch import func_on_device
+else:
+    raise ValueError(f"Unsupported backend: {keras.backend.backend()}")

zea/backend/jax/__init__.py

@@ -1,47 +1,21 @@
 """Jax utilities for zea."""
 
 import jax
-import numpy as np
 
 
-def on_device_jax(func, inputs, device, return_numpy=False, **kwargs):
-    """Applies a JAX function to inputs on a specified device.
-
+def str_to_jax_device(device):
+    """Convert a device string to a JAX device.
     Args:
-        func (callable): The function to apply.
-        inputs (ndarray): Input array.
-        device (str): Device string, e.g. ``'cuda'``, ``'gpu'``, or ``'cpu'``.
-        return_numpy (bool, optional): Whether to convert output
-            data back to numpy. Defaults to False.
-        **kwargs: Additional keyword arguments to be passed to the ``func``.
-
+        device (str): Device string, e.g. ``'gpu:0'``, or ``'cpu:0'``.
     Returns:
-        jax.numpy.DeviceArray or ndarray: The output data.
-
-    Raises:
-        AssertionError: If ``func`` is not a function from the JAX library.
-
-    Note:
-        This function converts the ``inputs`` array to a JAX array and moves
-        it to the specified ``device``. It then applies the ``func`` function to the inputs
-        and returns the output data. If the output is a dictionary, it extracts the first value
-        from the dictionary. If ``return_numpy`` is True, it converts the output data back to a
-        numpy array before returning.
-
-    Example:
-        .. code-block:: python
-
-            import jax.numpy as jnp
-
+        jax.Device: The corresponding JAX device.
+    """
 
-            def square(x):
-                return x**2
+    if not isinstance(device, str):
+        raise ValueError(f"Device must be a string, got {type(device)}")
 
+    device = device.lower().replace("cuda", "gpu")
 
-            inputs = [1, 2, 3, 4, 5]
-            device = "gpu"
-            output = on_device_jax(square, inputs, device)
-    """
     device = device.split(":")
     if len(device) == 2:
         device_type, device_number = device
@@ -51,20 +25,9 @@ def on_device_jax(func, inputs, device, return_numpy=False, **kwargs):
         device_type = device[0]
         device_number = 0
 
-    if device_number > len(jax.devices(device_type)):
-        raise ValueError(
-            f"Device {device} is not available from JAX devices: {jax.devices(device_type)}"
-        )
-
-    jax_device = jax.devices(device_type)[device_number]
-
-    with jax.default_device(jax_device):
-        outputs = func(inputs, **kwargs)
-
-    if isinstance(outputs, dict):
-        outputs = list(outputs.values())[0]
-
-    if return_numpy:
-        outputs = np.array(outputs)
-
-    return outputs
+    available = jax.devices(device_type)
+    if len(available) == 0:
+        raise ValueError(f"No JAX devices available for type '{device_type}'.")
+    if device_number < 0 or device_number >= len(available):
+        raise ValueError(f"Device '{device}' is not available; JAX devices found: {available}")
+    return available[device_number]

zea/backend/tensorflow/__init__.py

@@ -15,52 +15,3 @@ sys.path = [str(p) if isinstance(p, PosixPath) else p for p in sys.path]
 import tensorflow as tf  # noqa: E402
 
 from .dataloader import make_dataloader  # noqa: E402
-
-
-def on_device_tf(func, inputs, device, return_numpy=False, **kwargs):
-    """Applies a Tensorflow function to inputs on a specified device.
-
-    Args:
-        func (function): Function to apply to the input data.
-        inputs (ndarray): Input array.
-        device (str): Device string, e.g. ``'cuda'``, ``'gpu'``, or ``'cpu'``.
-        return_numpy (bool, optional): Whether to convert output
-            data back to numpy. Defaults to False.
-        **kwargs: Additional keyword arguments to be passed to the ``func``.
-
-    Returns:
-        tf.Tensor or ndarray: The output data.
-
-    Raises:
-        AssertionError: If ``func`` is not a function from the tensorflow library.
-
-    Note:
-        This function converts the ``inputs`` array to a tf.Tensor and moves
-        it to the specified ``device``. It then applies the ``func`` function to the inputs
-        and returns the output data. If the output is a dictionary, it extracts the first value
-        from the dictionary. If ``return_numpy`` is True, it converts the output data back to a
-        numpy array before returning.
-
-    Example:
-        .. code-block:: python
-
-            import tensorflow as tf
-
-
-            def square(x):
-                return x**2
-
-
-            inputs = [1, 2, 3, 4, 5]
-            device = "cuda"
-            output = on_device_tf(square, inputs, device)
-    """
-    device = device.replace("cuda", "gpu")
-
-    with tf.device(device):
-        outputs = func(inputs, **kwargs)
-
-    if return_numpy:
-        if not isinstance(outputs, np.ndarray):
-            outputs = outputs.numpy()
-    return outputs

zea/backend/torch/__init__.py

@@ -3,72 +3,37 @@
 Initialize modules for registries.
 """
 
-import numpy as np
 import torch
 
 
-def on_device_torch(func, inputs, device, return_numpy=False, **kwargs):
-    """Applies a Torch function to inputs on a specified device.
+def func_on_device(func, device, *args, **kwargs):
+    """Moves all tensor arguments of a function to a specified device before calling it.
 
     Args:
-        func (function): Function to apply to the input data.
-        inputs (ndarray): Input array.
-        device (str): Device string, e.g. ``'cuda'``, ``'gpu'``, or ``'cpu'``.
-        return_numpy (bool, optional): Whether to convert output
-            data back to numpy. Defaults to False.
-        **kwargs: Additional keyword arguments to be passed to the ``func``.
-
+        func (callable): Function to be called.
+        device (str or torch.device): Device to move tensors to.
+        *args: Positional arguments to be passed to the function.
+        **kwargs: Keyword arguments to be passed to the function.
     Returns:
-        torch.Tensor or ndarray: The output data.
-
-    Raises:
-        AssertionError: If ``func`` is not a function from the torch library.
-
-    Note:
-        This function converts the ``inputs`` array to a torch.Tensor and moves it to
-        the specified ``device``. It then applies the ``func`` function to the inputs and
-        returns the output data. If the output is a dictionary, it extracts the first value
-        from the dictionary. If ``return_numpy`` is True, it converts the output data back to a
-        numpy array before returning.
-
-    Example:
-        .. code-block:: python
-
-            import torch
-
-
-            def square(x):
-                return x**2
-
-
-            inputs = [1, 2, 3, 4, 5]
-            device = "cuda"
-            output = on_device_torch(square, inputs, device)
-            print(output)
+        The output of the function.
     """
-    device = device.replace("gpu", "cuda")
-
-    if not isinstance(inputs, torch.Tensor):
-        inputs = torch.tensor(inputs)
-
-    inputs = inputs.to(device)
-
-    # check that function is a function from torch library
-    # assert "torch" in str(type(func)), f"func: {func} should be a torch function"
-    if hasattr(func, "to"):
-        func = func.to(device)
-
-    with torch.device(device):
-        outputs = func(inputs, **kwargs)
-
-    if isinstance(outputs, dict):
-        # depends a bit how flexible we want to be...
-        # but for now quick and dirty solution
-        key = list(outputs.keys())[0]
-        outputs = outputs[key]
-
-    if return_numpy:
-        if not isinstance(outputs, np.ndarray):
-            outputs = outputs.cpu().numpy()
-
-    return outputs
+    if device is None:
+        return func(*args, **kwargs)
+
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    def move_to_device(x):
+        if isinstance(x, torch.Tensor):
+            return x.to(device)
+        elif isinstance(x, (list, tuple)):
+            return type(x)(move_to_device(i) for i in x)
+        elif isinstance(x, dict):
+            return {k: move_to_device(v) for k, v in x.items()}
+        else:
+            return x
+
+    args = move_to_device(args)
+    kwargs = move_to_device(kwargs)
+
+    return func(*args, **kwargs)

zea/data/layers.py

@@ -25,7 +25,7 @@ class Resizer(TFDataLayer):
     Resize layer for resizing images. Can deal with N-dimensional images.
     Can do resize, center_crop, random_crop and crop_or_pad.
 
-    Can be used in tf.data pipelines.
+    Can be used in `tf.data` pipelines.
     """
 
     def __init__(
@@ -36,7 +36,6 @@ class Resizer(TFDataLayer):
         seed: int | None = None,
         **resize_kwargs,
     ):
-        # noqa: E501
         """
         Initializes the data loader with the specified parameters.
 
@@ -47,7 +46,7 @@ class Resizer(TFDataLayer):
                 ['random_crop'](https://keras.io/api/layers/preprocessing_layers/image_augmentation/random_crop/),
                 ['resize'](https://keras.io/api/layers/preprocessing_layers/image_preprocessing/resizing/),
                 'crop_or_pad': resizes an image to a target width and height by either centrally
-                    cropping the image, padding it evenly with zeros or a combination of both.
+                cropping the image, padding it evenly with zeros or a combination of both.
             resize_axes (tuple | None, optional): The axes along which to resize.
                 Must be of length 2. Defaults to None. In that case, can only process
                 default tensors of shape (batch, height, width, channels), where the

zea/internal/registry.py

@@ -200,7 +200,7 @@ tf_beamformer_registry = RegisterDecorator(items_to_register=["name", "framework
 
 torch_beamformer_registry = RegisterDecorator(items_to_register=["name", "framework"])
 
-metrics_registry = RegisterDecorator(items_to_register=["name", "framework", "supervised"])
+metrics_registry = RegisterDecorator(items_to_register=["name", "paired"])
 
 checks_registry = RegisterDecorator(items_to_register=["data_type"])
 ops_registry = RegisterDecorator(items_to_register=["name"])