explainiverse 0.7.0__tar.gz → 0.8.0__tar.gz

{explainiverse-0.7.0 → explainiverse-0.8.0}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.1
 Name: explainiverse
-Version: 0.7.0
-Summary: Unified, extensible explainability framework supporting LIME, SHAP, Anchors, Counterfactuals, PDP, ALE, SAGE, and more
+Version: 0.8.0
+Summary: Unified, extensible explainability framework supporting 18 XAI methods including LIME, SHAP, LRP, TCAV, GradCAM, and more
 Home-page: https://github.com/jemsbhai/explainiverse
 License: MIT
 Keywords: xai,explainability,interpretability,machine-learning,lime,shap,anchors
@@ -35,7 +35,7 @@ Description-Content-Type: text/markdown
 [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 
-**Explainiverse** is a unified, extensible Python framework for Explainable AI (XAI). It provides a standardized interface for **17 state-of-the-art explanation methods** across local, global, gradient-based, concept-based, and example-based paradigms, along with **comprehensive evaluation metrics** for assessing explanation quality.
+**Explainiverse** is a unified, extensible Python framework for Explainable AI (XAI). It provides a standardized interface for **18 state-of-the-art explanation methods** across local, global, gradient-based, concept-based, and example-based paradigms, along with **comprehensive evaluation metrics** for assessing explanation quality.
 
 ---
 
@@ -43,7 +43,7 @@ Description-Content-Type: text/markdown
 
 | Feature | Description |
 |---------|-------------|
-| **17 Explainers** | LIME, KernelSHAP, TreeSHAP, Integrated Gradients, DeepLIFT, DeepSHAP, SmoothGrad, Saliency Maps, GradCAM/GradCAM++, TCAV, Anchors, Counterfactual, Permutation Importance, PDP, ALE, SAGE, ProtoDash |
+| **18 Explainers** | LIME, KernelSHAP, TreeSHAP, Integrated Gradients, DeepLIFT, DeepSHAP, SmoothGrad, Saliency Maps, GradCAM/GradCAM++, LRP, TCAV, Anchors, Counterfactual, Permutation Importance, PDP, ALE, SAGE, ProtoDash |
 | **8 Evaluation Metrics** | Faithfulness (PGI, PGU, Comprehensiveness, Sufficiency, Correlation) and Stability (RIS, ROS, Lipschitz) |
 | **Unified API** | Consistent `BaseExplainer` interface with standardized `Explanation` output |
 | **Plugin Registry** | Filter explainers by scope, model type, data type; automatic recommendations |
@@ -66,6 +66,7 @@ Description-Content-Type: text/markdown
 | **SmoothGrad** | Gradient | [Smilkov et al., 2017](https://arxiv.org/abs/1706.03825) |
 | **Saliency Maps** | Gradient | [Simonyan et al., 2014](https://arxiv.org/abs/1312.6034) |
 | **GradCAM / GradCAM++** | Gradient (CNN) | [Selvaraju et al., 2017](https://arxiv.org/abs/1610.02391) |
+| **LRP** | Decomposition | [Bach et al., 2015](https://doi.org/10.1371/journal.pone.0130140) |
 | **TCAV** | Concept-Based | [Kim et al., 2018](https://arxiv.org/abs/1711.11279) |
 | **Anchors** | Rule-Based | [Ribeiro et al., 2018](https://ojs.aaai.org/index.php/AAAI/article/view/11491) |
 | **Counterfactual** | Contrastive | [Mothilal et al., 2020](https://arxiv.org/abs/1905.07697) |
@@ -143,7 +144,7 @@ adapter = SklearnAdapter(model, class_names=iris.target_names.tolist())
 # List all available explainers
 print(default_registry.list_explainers())
 # ['lime', 'shap', 'treeshap', 'integrated_gradients', 'deeplift', 'deepshap', 
-#  'smoothgrad', 'saliency', 'gradcam', 'tcav', 'anchors', 'counterfactual', 
+#  'smoothgrad', 'saliency', 'gradcam', 'lrp', 'tcav', 'anchors', 'counterfactual', 
 #  'protodash', 'permutation_importance', 'partial_dependence', 'ale', 'sage']
 
 # Create an explainer via registry
@@ -211,6 +212,70 @@ print(f"Attributions: {explanation.explanation_data['feature_attributions']}")
 print(f"Convergence δ: {explanation.explanation_data['convergence_delta']:.6f}")
 ```
 
+### Layer-wise Relevance Propagation (LRP)
+
+```python
+from explainiverse.explainers.gradient import LRPExplainer
+
+# LRP - Decomposition-based attribution with conservation property
+explainer = LRPExplainer(
+    model=adapter,
+    feature_names=feature_names,
+    class_names=class_names,
+    rule="epsilon",       # Propagation rule: epsilon, gamma, alpha_beta, z_plus, composite
+    epsilon=1e-6          # Stabilization constant
+)
+
+# Basic explanation
+explanation = explainer.explain(X[0], target_class=0)
+print(explanation.explanation_data["feature_attributions"])
+
+# Verify conservation property (sum of attributions ≈ target output)
+explanation = explainer.explain(X[0], return_convergence_delta=True)
+print(f"Conservation delta: {explanation.explanation_data['convergence_delta']:.6f}")
+
+# Compare different LRP rules
+comparison = explainer.compare_rules(X[0], rules=["epsilon", "gamma", "z_plus"])
+for rule, result in comparison.items():
+    print(f"{rule}: top feature = {result['top_feature']}")
+
+# Layer-wise relevance analysis
+layer_result = explainer.explain_with_layer_relevances(X[0])
+for layer, relevances in layer_result["layer_relevances"].items():
+    print(f"{layer}: sum = {sum(relevances):.4f}")
+
+# Composite rules: different rules for different layers
+explainer_composite = LRPExplainer(
+    model=adapter,
+    feature_names=feature_names,
+    class_names=class_names,
+    rule="composite"
+)
+explainer_composite.set_composite_rule({
+    0: "z_plus",    # Input layer: focus on what's present
+    2: "epsilon",   # Middle layers: balanced
+    4: "epsilon"    # Output layer
+})
+explanation = explainer_composite.explain(X[0])
+```
+
+**LRP Propagation Rules:**
+
+| Rule | Description | Use Case |
+|------|-------------|----------|
+| `epsilon` | Adds stabilization constant | General purpose (default) |
+| `gamma` | Enhances positive contributions | Image classification |
+| `alpha_beta` | Separates pos/neg (α-β=1) | Fine-grained control |
+| `z_plus` | Only positive weights | Input layers, what's present |
+| `composite` | Different rules per layer | Best practice for deep nets |
+
+**Supported Layers:**
+- Linear, Conv2d
+- BatchNorm1d, BatchNorm2d
+- ReLU, LeakyReLU, ELU, Tanh, Sigmoid, GELU
+- MaxPool2d, AvgPool2d, AdaptiveAvgPool2d
+- Flatten, Dropout
+
 ### DeepLIFT and DeepSHAP
 
 ```python
@@ -602,7 +667,7 @@ explainiverse/
 │   └── pytorch_adapter.py  # With gradient support
 ├── explainers/
 │   ├── attribution/      # LIME, SHAP, TreeSHAP
-│   ├── gradient/         # IG, DeepLIFT, DeepSHAP, SmoothGrad, Saliency, GradCAM, TCAV
+│   ├── gradient/         # IG, DeepLIFT, DeepSHAP, SmoothGrad, Saliency, GradCAM, LRP, TCAV
 │   ├── rule_based/       # Anchors
 │   ├── counterfactual/   # DiCE-style
 │   ├── global_explainers/  # Permutation, PDP, ALE, SAGE
@@ -626,10 +691,10 @@ poetry run pytest
 poetry run pytest --cov=explainiverse --cov-report=html
 
 # Run specific test file
-poetry run pytest tests/test_smoothgrad.py -v
+poetry run pytest tests/test_lrp.py -v
 
 # Run specific test class
-poetry run pytest tests/test_smoothgrad.py::TestSmoothGradBasic -v
+poetry run pytest tests/test_lrp.py::TestLRPConv2d -v
 ```
 
 ---
@@ -640,6 +705,7 @@ poetry run pytest tests/test_smoothgrad.py::TestSmoothGradBasic -v
 - [x] Core framework (BaseExplainer, Explanation, Registry)
 - [x] Perturbation methods: LIME, KernelSHAP, TreeSHAP
 - [x] Gradient methods: Integrated Gradients, DeepLIFT, DeepSHAP, SmoothGrad, Saliency Maps, GradCAM/GradCAM++
+- [x] Decomposition methods: Layer-wise Relevance Propagation (LRP) with ε, γ, αβ, z⁺, composite rules
 - [x] Concept-based: TCAV (Testing with Concept Activation Vectors)
 - [x] Rule-based: Anchors
 - [x] Counterfactual: DiCE-style
@@ -649,9 +715,6 @@ poetry run pytest tests/test_smoothgrad.py::TestSmoothGradBasic -v
 - [x] Evaluation: Stability metrics (RIS, ROS, Lipschitz)
 - [x] PyTorch adapter with gradient support
 
-### In Progress 🚧
-- [ ] Layer-wise Relevance Propagation (LRP)
-
 ### Planned 📋
 - [ ] Attention-based explanations (for Transformers)
 - [ ] TensorFlow/Keras adapter
@@ -671,7 +734,7 @@ If you use Explainiverse in your research, please cite:
   author = {Syed, Muntaser},
   year = {2025},
   url = {https://github.com/jemsbhai/explainiverse},
-  version = {0.7.0}
+  version = {0.8.0}
 }
 ```
 
@@ -699,5 +762,5 @@ MIT License - see [LICENSE](LICENSE) for details.
 
 ## Acknowledgments
 
-Explainiverse builds upon the foundational work of many researchers in the XAI community. We thank the authors of LIME, SHAP, Integrated Gradients, DeepLIFT, GradCAM, TCAV, Anchors, DiCE, ALE, SAGE, and ProtoDash for their contributions to interpretable machine learning.
+Explainiverse builds upon the foundational work of many researchers in the XAI community. We thank the authors of LIME, SHAP, Integrated Gradients, DeepLIFT, LRP, GradCAM, TCAV, Anchors, DiCE, ALE, SAGE, and ProtoDash for their contributions to interpretable machine learning.
 

{explainiverse-0.7.0 → explainiverse-0.8.0}/README.md

@@ -4,7 +4,7 @@
 [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 
-**Explainiverse** is a unified, extensible Python framework for Explainable AI (XAI). It provides a standardized interface for **17 state-of-the-art explanation methods** across local, global, gradient-based, concept-based, and example-based paradigms, along with **comprehensive evaluation metrics** for assessing explanation quality.
+**Explainiverse** is a unified, extensible Python framework for Explainable AI (XAI). It provides a standardized interface for **18 state-of-the-art explanation methods** across local, global, gradient-based, concept-based, and example-based paradigms, along with **comprehensive evaluation metrics** for assessing explanation quality.
 
 ---
 
@@ -12,7 +12,7 @@
 
 | Feature | Description |
 |---------|-------------|
-| **17 Explainers** | LIME, KernelSHAP, TreeSHAP, Integrated Gradients, DeepLIFT, DeepSHAP, SmoothGrad, Saliency Maps, GradCAM/GradCAM++, TCAV, Anchors, Counterfactual, Permutation Importance, PDP, ALE, SAGE, ProtoDash |
+| **18 Explainers** | LIME, KernelSHAP, TreeSHAP, Integrated Gradients, DeepLIFT, DeepSHAP, SmoothGrad, Saliency Maps, GradCAM/GradCAM++, LRP, TCAV, Anchors, Counterfactual, Permutation Importance, PDP, ALE, SAGE, ProtoDash |
 | **8 Evaluation Metrics** | Faithfulness (PGI, PGU, Comprehensiveness, Sufficiency, Correlation) and Stability (RIS, ROS, Lipschitz) |
 | **Unified API** | Consistent `BaseExplainer` interface with standardized `Explanation` output |
 | **Plugin Registry** | Filter explainers by scope, model type, data type; automatic recommendations |
@@ -35,6 +35,7 @@
 | **SmoothGrad** | Gradient | [Smilkov et al., 2017](https://arxiv.org/abs/1706.03825) |
 | **Saliency Maps** | Gradient | [Simonyan et al., 2014](https://arxiv.org/abs/1312.6034) |
 | **GradCAM / GradCAM++** | Gradient (CNN) | [Selvaraju et al., 2017](https://arxiv.org/abs/1610.02391) |
+| **LRP** | Decomposition | [Bach et al., 2015](https://doi.org/10.1371/journal.pone.0130140) |
 | **TCAV** | Concept-Based | [Kim et al., 2018](https://arxiv.org/abs/1711.11279) |
 | **Anchors** | Rule-Based | [Ribeiro et al., 2018](https://ojs.aaai.org/index.php/AAAI/article/view/11491) |
 | **Counterfactual** | Contrastive | [Mothilal et al., 2020](https://arxiv.org/abs/1905.07697) |
@@ -112,7 +113,7 @@ adapter = SklearnAdapter(model, class_names=iris.target_names.tolist())
 # List all available explainers
 print(default_registry.list_explainers())
 # ['lime', 'shap', 'treeshap', 'integrated_gradients', 'deeplift', 'deepshap', 
-#  'smoothgrad', 'saliency', 'gradcam', 'tcav', 'anchors', 'counterfactual', 
+#  'smoothgrad', 'saliency', 'gradcam', 'lrp', 'tcav', 'anchors', 'counterfactual', 
 #  'protodash', 'permutation_importance', 'partial_dependence', 'ale', 'sage']
 
 # Create an explainer via registry
@@ -180,6 +181,70 @@ print(f"Attributions: {explanation.explanation_data['feature_attributions']}")
 print(f"Convergence δ: {explanation.explanation_data['convergence_delta']:.6f}")
 ```
 
+### Layer-wise Relevance Propagation (LRP)
+
+```python
+from explainiverse.explainers.gradient import LRPExplainer
+
+# LRP - Decomposition-based attribution with conservation property
+explainer = LRPExplainer(
+    model=adapter,
+    feature_names=feature_names,
+    class_names=class_names,
+    rule="epsilon",       # Propagation rule: epsilon, gamma, alpha_beta, z_plus, composite
+    epsilon=1e-6          # Stabilization constant
+)
+
+# Basic explanation
+explanation = explainer.explain(X[0], target_class=0)
+print(explanation.explanation_data["feature_attributions"])
+
+# Verify conservation property (sum of attributions ≈ target output)
+explanation = explainer.explain(X[0], return_convergence_delta=True)
+print(f"Conservation delta: {explanation.explanation_data['convergence_delta']:.6f}")
+
+# Compare different LRP rules
+comparison = explainer.compare_rules(X[0], rules=["epsilon", "gamma", "z_plus"])
+for rule, result in comparison.items():
+    print(f"{rule}: top feature = {result['top_feature']}")
+
+# Layer-wise relevance analysis
+layer_result = explainer.explain_with_layer_relevances(X[0])
+for layer, relevances in layer_result["layer_relevances"].items():
+    print(f"{layer}: sum = {sum(relevances):.4f}")
+
+# Composite rules: different rules for different layers
+explainer_composite = LRPExplainer(
+    model=adapter,
+    feature_names=feature_names,
+    class_names=class_names,
+    rule="composite"
+)
+explainer_composite.set_composite_rule({
+    0: "z_plus",    # Input layer: focus on what's present
+    2: "epsilon",   # Middle layers: balanced
+    4: "epsilon"    # Output layer
+})
+explanation = explainer_composite.explain(X[0])
+```
+
+**LRP Propagation Rules:**
+
+| Rule | Description | Use Case |
+|------|-------------|----------|
+| `epsilon` | Adds stabilization constant | General purpose (default) |
+| `gamma` | Enhances positive contributions | Image classification |
+| `alpha_beta` | Separates pos/neg (α-β=1) | Fine-grained control |
+| `z_plus` | Only positive weights | Input layers, what's present |
+| `composite` | Different rules per layer | Best practice for deep nets |
+
+**Supported Layers:**
+- Linear, Conv2d
+- BatchNorm1d, BatchNorm2d
+- ReLU, LeakyReLU, ELU, Tanh, Sigmoid, GELU
+- MaxPool2d, AvgPool2d, AdaptiveAvgPool2d
+- Flatten, Dropout
+
 ### DeepLIFT and DeepSHAP
 
 ```python
@@ -571,7 +636,7 @@ explainiverse/
 │   └── pytorch_adapter.py  # With gradient support
 ├── explainers/
 │   ├── attribution/      # LIME, SHAP, TreeSHAP
-│   ├── gradient/         # IG, DeepLIFT, DeepSHAP, SmoothGrad, Saliency, GradCAM, TCAV
+│   ├── gradient/         # IG, DeepLIFT, DeepSHAP, SmoothGrad, Saliency, GradCAM, LRP, TCAV
 │   ├── rule_based/       # Anchors
 │   ├── counterfactual/   # DiCE-style
 │   ├── global_explainers/  # Permutation, PDP, ALE, SAGE
@@ -595,10 +660,10 @@ poetry run pytest
 poetry run pytest --cov=explainiverse --cov-report=html
 
 # Run specific test file
-poetry run pytest tests/test_smoothgrad.py -v
+poetry run pytest tests/test_lrp.py -v
 
 # Run specific test class
-poetry run pytest tests/test_smoothgrad.py::TestSmoothGradBasic -v
+poetry run pytest tests/test_lrp.py::TestLRPConv2d -v
 ```
 
 ---
@@ -609,6 +674,7 @@ poetry run pytest tests/test_smoothgrad.py::TestSmoothGradBasic -v
 - [x] Core framework (BaseExplainer, Explanation, Registry)
 - [x] Perturbation methods: LIME, KernelSHAP, TreeSHAP
 - [x] Gradient methods: Integrated Gradients, DeepLIFT, DeepSHAP, SmoothGrad, Saliency Maps, GradCAM/GradCAM++
+- [x] Decomposition methods: Layer-wise Relevance Propagation (LRP) with ε, γ, αβ, z⁺, composite rules
 - [x] Concept-based: TCAV (Testing with Concept Activation Vectors)
 - [x] Rule-based: Anchors
 - [x] Counterfactual: DiCE-style
@@ -618,9 +684,6 @@ poetry run pytest tests/test_smoothgrad.py::TestSmoothGradBasic -v
 - [x] Evaluation: Stability metrics (RIS, ROS, Lipschitz)
 - [x] PyTorch adapter with gradient support
 
-### In Progress 🚧
-- [ ] Layer-wise Relevance Propagation (LRP)
-
 ### Planned 📋
 - [ ] Attention-based explanations (for Transformers)
 - [ ] TensorFlow/Keras adapter
@@ -640,7 +703,7 @@ If you use Explainiverse in your research, please cite:
   author = {Syed, Muntaser},
   year = {2025},
   url = {https://github.com/jemsbhai/explainiverse},
-  version = {0.7.0}
+  version = {0.8.0}
 }
 ```
 
@@ -668,4 +731,4 @@ MIT License - see [LICENSE](LICENSE) for details.
 
 ## Acknowledgments
 
-Explainiverse builds upon the foundational work of many researchers in the XAI community. We thank the authors of LIME, SHAP, Integrated Gradients, DeepLIFT, GradCAM, TCAV, Anchors, DiCE, ALE, SAGE, and ProtoDash for their contributions to interpretable machine learning.
+Explainiverse builds upon the foundational work of many researchers in the XAI community. We thank the authors of LIME, SHAP, Integrated Gradients, DeepLIFT, LRP, GradCAM, TCAV, Anchors, DiCE, ALE, SAGE, and ProtoDash for their contributions to interpretable machine learning.

{explainiverse-0.7.0 → explainiverse-0.8.0}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "explainiverse"
-version = "0.7.0"
-description = "Unified, extensible explainability framework supporting LIME, SHAP, Anchors, Counterfactuals, PDP, ALE, SAGE, and more"
+version = "0.8.0"
+description = "Unified, extensible explainability framework supporting 18 XAI methods including LIME, SHAP, LRP, TCAV, GradCAM, and more"
 authors = ["Muntaser Syed <jemsbhai@gmail.com>"]
 license = "MIT"
 readme = "README.md"

{explainiverse-0.7.0 → explainiverse-0.8.0}/src/explainiverse/__init__.py

@@ -2,9 +2,10 @@
 """
 Explainiverse - A unified, extensible explainability framework.
 
-Supports multiple XAI methods including LIME, SHAP, TreeSHAP, Anchors, 
-Counterfactuals, Permutation Importance, PDP, ALE, and SAGE through a 
-consistent interface.
+Supports 18 state-of-the-art XAI methods including LIME, SHAP, TreeSHAP,
+Integrated Gradients, DeepLIFT, DeepSHAP, LRP, GradCAM, TCAV, Anchors,
+Counterfactuals, Permutation Importance, PDP, ALE, SAGE, and ProtoDash
+through a consistent interface.
 
 Quick Start:
     from explainiverse import default_registry
@@ -33,7 +34,7 @@ from explainiverse.adapters.sklearn_adapter import SklearnAdapter
 from explainiverse.adapters import TORCH_AVAILABLE
 from explainiverse.engine.suite import ExplanationSuite
 
-__version__ = "0.7.0"
+__version__ = "0.8.0"
 
 __all__ = [
     # Core

{explainiverse-0.7.0 → explainiverse-0.8.0}/src/explainiverse/adapters/pytorch_adapter.py

@@ -25,7 +25,7 @@ Example:
 """
 
 import numpy as np
-from typing import List, Optional, Union, Callable
+from typing import List, Optional, Union, Tuple
 
 from .base_adapter import BaseModelAdapter
 
@@ -57,6 +57,11 @@ class PyTorchAdapter(BaseModelAdapter):
     explainability methods. Handles device management, tensor/numpy
     conversions, and supports both classification and regression tasks.
     
+    Supports:
+        - Multi-class classification (output shape: [batch, n_classes])
+        - Binary classification (output shape: [batch, 1] or [batch])
+        - Regression (output shape: [batch, n_outputs] or [batch])
+    
     Attributes:
         model: The PyTorch model (nn.Module)
         task: "classification" or "regression"
@@ -150,11 +155,27 @@ class PyTorchAdapter(BaseModelAdapter):
     def _apply_activation(self, output: "torch.Tensor") -> "torch.Tensor":
         """Apply output activation function."""
         if self.output_activation == "softmax":
+            # Handle different output shapes
+            if output.dim() == 1 or (output.dim() == 2 and output.shape[1] == 1):
+                # Binary: apply sigmoid instead of softmax
+                return torch.sigmoid(output)
             return torch.softmax(output, dim=-1)
         elif self.output_activation == "sigmoid":
             return torch.sigmoid(output)
         return output
     
+    def _normalize_output_shape(self, output: "torch.Tensor") -> "torch.Tensor":
+        """
+        Normalize output to consistent 2D shape (batch, outputs).
+        
+        Handles:
+            - (batch,) -> (batch, 1)
+            - (batch, n) -> (batch, n)
+        """
+        if output.dim() == 1:
+            return output.unsqueeze(-1)
+        return output
+    
     def predict(self, data: np.ndarray) -> np.ndarray:
         """
         Generate predictions for input data.
@@ -183,16 +204,66 @@ class PyTorchAdapter(BaseModelAdapter):
                 tensor_batch = self._to_tensor(batch)
                 
                 output = self.model(tensor_batch)
+                output = self._normalize_output_shape(output)
                 output = self._apply_activation(output)
                 outputs.append(self._to_numpy(output))
         
         return np.vstack(outputs)
     
+    def _get_target_scores(
+        self,
+        output: "torch.Tensor",
+        target_class: Optional[Union[int, "torch.Tensor"]] = None
+    ) -> "torch.Tensor":
+        """
+        Extract target scores for gradient computation.
+        
+        Handles both multi-class and binary classification outputs.
+        
+        Args:
+            output: Raw model output (logits)
+            target_class: Target class index or tensor of indices
+            
+        Returns:
+            Target scores tensor for backpropagation
+        """
+        batch_size = output.shape[0]
+        
+        # Normalize to 2D
+        if output.dim() == 1:
+            output = output.unsqueeze(-1)
+        
+        n_outputs = output.shape[1]
+        
+        if self.task == "classification":
+            if n_outputs == 1:
+                # Binary classification with single logit
+                # Score is the logit itself (positive class score)
+                return output.squeeze(-1)
+            else:
+                # Multi-class classification
+                if target_class is None:
+                    target_class = output.argmax(dim=-1)
+                elif isinstance(target_class, int):
+                    target_class = torch.tensor(
+                        [target_class] * batch_size,
+                        device=self.device
+                    )
+                
+                # Gather scores for target class
+                return output.gather(1, target_class.view(-1, 1)).squeeze(-1)
+        else:
+            # Regression: use first output or sum of outputs
+            if n_outputs == 1:
+                return output.squeeze(-1)
+            else:
+                return output.sum(dim=-1)
+    
     def predict_with_gradients(
         self,
         data: np.ndarray,
         target_class: Optional[int] = None
-    ) -> tuple:
+    ) -> Tuple[np.ndarray, np.ndarray]:
         """
         Generate predictions and compute gradients w.r.t. inputs.
         
@@ -203,11 +274,17 @@ class PyTorchAdapter(BaseModelAdapter):
             data: Input data as numpy array.
             target_class: Class index for gradient computation.
                          If None, uses the predicted class.
+                         For binary classification with single output,
+                         this is ignored (gradient w.r.t. the single logit).
         
         Returns:
             Tuple of (predictions, gradients) as numpy arrays.
+            - predictions: (batch, n_classes) probabilities
+            - gradients: same shape as input data
         """
         data = np.array(data)
+        original_shape = data.shape
+        
         if data.ndim == 1:
             data = data.reshape(1, -1)
         
@@ -217,20 +294,13 @@ class PyTorchAdapter(BaseModelAdapter):
         
         # Forward pass
         output = self.model(tensor_data)
-        activated_output = self._apply_activation(output)
         
-        # Determine target for gradient
-        if self.task == "classification":
-            if target_class is None:
-                target_class = output.argmax(dim=-1)
-            elif isinstance(target_class, int):
-                target_class = torch.tensor([target_class] * data.shape[0], device=self.device)
-            
-            # Select target class scores for gradient
-            target_scores = output.gather(1, target_class.view(-1, 1)).squeeze()
-        else:
-            # Regression: gradient w.r.t. output
-            target_scores = output.squeeze()
+        # Get activated output for return
+        output_normalized = self._normalize_output_shape(output)
+        activated_output = self._apply_activation(output_normalized)
+        
+        # Get target scores for gradient computation
+        target_scores = self._get_target_scores(output, target_class)
         
         # Backward pass
         if target_scores.dim() == 0:
@@ -295,7 +365,7 @@ class PyTorchAdapter(BaseModelAdapter):
         data: np.ndarray,
         layer_name: str,
         target_class: Optional[int] = None
-    ) -> tuple:
+    ) -> Tuple[np.ndarray, np.ndarray]:
         """
         Get gradients of output w.r.t. a specific layer's activations.
         
@@ -339,15 +409,8 @@ class PyTorchAdapter(BaseModelAdapter):
             
             output = self.model(tensor_data)
             
-            if self.task == "classification":
-                if target_class is None:
-                    target_class = output.argmax(dim=-1)
-                elif isinstance(target_class, int):
-                    target_class = torch.tensor([target_class] * data.shape[0], device=self.device)
-                
-                target_scores = output.gather(1, target_class.view(-1, 1)).squeeze()
-            else:
-                target_scores = output.squeeze()
+            # Get target scores using the new method
+            target_scores = self._get_target_scores(output, target_class)
             
             if target_scores.dim() == 0:
                 target_scores.backward()