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

zea/__init__.py

@@ -7,36 +7,71 @@ from . import log
 
 # dynamically add __version__ attribute (see pyproject.toml)
 # __version__ = __import__("importlib.metadata").metadata.version(__package__)
-__version__ = "0.0.5"
+__version__ = "0.0.7"
 
 
 def _bootstrap_backend():
     """Setup function to initialize the zea package."""
 
     def _check_backend_installed():
-        """Assert that at least one ML backend (torch, tensorflow, jax) is installed.
-        If not, raise an AssertionError with a helpful install message.
-        """
-
-        ml_backends = ["torch", "tensorflow", "jax"]
-        for backend in ml_backends:
-            if importlib.util.find_spec(backend) is not None:
-                return
+        """Verify that the required ML backend is installed.
 
-        backend_env = os.environ.get("KERAS_BACKEND", "numpy")
-        install_guide_urls = {
+        Raises ImportError if:
+        1. No ML backend (torch, tensorflow, jax) is installed
+        2. KERAS_BACKEND points to a backend that is not installed
+        """
+        ML_BACKENDS = ["torch", "tensorflow", "jax"]
+        INSTALL_URLS = {
             "torch": "https://pytorch.org/get-started/locally/",
             "tensorflow": "https://www.tensorflow.org/install",
             "jax": "https://docs.jax.dev/en/latest/installation.html",
         }
-        guide_url = install_guide_urls.get(backend_env, "https://keras.io/getting_started/")
-        raise ImportError(
-            "No ML backend (torch, tensorflow, jax) installed in current environment. "
-            f"Please install at least one ML backend before importing {__package__} or "
-            f"any other library. Current KERAS_BACKEND is set to '{backend_env}', "
-            f"please install it first, see: {guide_url}. One simple alternative is to "
-            f"install with default backend: `pip install {__package__}[jax]`."
-        )
+        KERAS_DEFAULT_BACKEND = "tensorflow"
+        DOCS_URL = "https://zea.readthedocs.io/en/latest/installation.html"
+
+        # Determine which backend Keras will try to use
+        backend_env = os.environ.get("KERAS_BACKEND")
+        effective_backend = backend_env or KERAS_DEFAULT_BACKEND
+
+        # Find all installed ML backends
+        installed_backends = [
+            backend for backend in ML_BACKENDS if importlib.util.find_spec(backend) is not None
+        ]
+
+        # Error if no backends are installed
+        if not installed_backends:
+            if backend_env:
+                backend_status = f"KERAS_BACKEND is set to '{backend_env}'"
+            else:
+                backend_status = f"KERAS_BACKEND is not set (defaults to '{KERAS_DEFAULT_BACKEND}')"
+            install_url = INSTALL_URLS.get(effective_backend, "https://keras.io/getting_started/")
+            raise ImportError(
+                f"No ML backend (torch, tensorflow, jax) installed in current "
+                f"environment. Please install at least one ML backend before importing "
+                f"{__package__}. {backend_status}, please install it first, see: "
+                f"{install_url}. One simple alternative is to install with default "
+                f"backend: `pip install {__package__}[jax]`. For more information, "
+                f"see: {DOCS_URL}"
+            )
+
+        # Error if the effective backend is not installed
+        # (skip numpy which doesn't need installation)
+        if effective_backend not in ["numpy"] and effective_backend not in installed_backends:
+            if backend_env:
+                backend_status = f"KERAS_BACKEND environment variable is set to '{backend_env}'"
+            else:
+                backend_status = (
+                    f"KERAS_BACKEND is not set, which defaults to '{KERAS_DEFAULT_BACKEND}'"
+                )
+            install_url = INSTALL_URLS.get(effective_backend, "https://keras.io/getting_started/")
+            raise ImportError(
+                f"{backend_status}, but this backend is not installed. "
+                f"Installed backends: {', '.join(installed_backends)}. "
+                f"Please either install '{effective_backend}' (see: {install_url}) "
+                f"or set KERAS_BACKEND to one of the installed backends "
+                f"(e.g., export KERAS_BACKEND={installed_backends[0]}). "
+                f"For more information, see: {DOCS_URL}"
+            )
 
     _check_backend_installed()
 

zea/agent/__init__.py

@@ -7,21 +7,21 @@ For a practical example, see :doc:`../notebooks/agent/agent_example`.
 Example usage
 ^^^^^^^^^^^^^
 
-.. code-block:: python
+.. doctest::
 
-    import zea
-    import numpy as np
+    >>> import zea
+    >>> import numpy as np
 
-    agent = zea.agent.selection.GreedyEntropy(
-        n_actions=7,
-        n_possible_actions=112,
-        img_width=112,
-        img_height=112,
-    )
+    >>> agent = zea.agent.selection.GreedyEntropy(
+    ...     n_actions=7,
+    ...     n_possible_actions=112,
+    ...     img_width=112,
+    ...     img_height=112,
+    ... )
 
-    # (batch, samples, height, width)
-    particles = np.random.rand(1, 10, 112, 112)
-    lines, mask = agent.sample(particles)
+    >>> # (batch, samples, height, width)
+    >>> particles = np.random.rand(1, 10, 112, 112)
+    >>> lines, mask = agent.sample(particles)  # doctest: +SKIP
 """
 
 from . import masks, selection

zea/agent/masks.py

@@ -91,7 +91,8 @@ def random_uniform_lines(
 
 def _assert_equal_spacing(n_actions, n_possible_actions):
     assert n_possible_actions % n_actions == 0, (
-        "Number of actions must divide evenly into possible actions to use equispaced sampling."
+        "Number of actions must divide evenly into possible actions to use equispaced sampling. "
+        "If you do not care about equal spacing, set `assert_equal_spacing=False`."
     )
 
 

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/tensorflow/dataloader.py

@@ -12,7 +12,8 @@ from keras.src.trainers.data_adapters import TFDatasetAdapter
 
 from zea.data.dataloader import H5Generator
 from zea.data.layers import Resizer
-from zea.utils import find_methods_with_return_type, translate
+from zea.internal.utils import find_methods_with_return_type
+from zea.tensor_ops import translate
 
 METHODS_THAT_RETURN_DATASET = find_methods_with_return_type(tf.data.Dataset, "DatasetV2")
 

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)