From 794dea5e138edf205ca1d3511d7d49017b6417c7 Mon Sep 17 00:00:00 2001
From: Kyrian Obikwelu <koshnawaza@gmail.com>
Date: Thu, 11 Apr 2024 01:23:57 +0100
Subject: [PATCH] Add support for tensor softmax for 2D tensors

---
 .../ZeroShotImageClassificationPipeline.php     |  5 ++---
 src/Processors/Processor.php                    |  7 ++-----
 src/Utils/Helpers.php                           |  2 --
 src/Utils/Tensor.php                            | 17 +++++++++++++++--
 4 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/src/Pipelines/ZeroShotImageClassificationPipeline.php b/src/Pipelines/ZeroShotImageClassificationPipeline.php
index 6cf0198..233f540 100644
--- a/src/Pipelines/ZeroShotImageClassificationPipeline.php
+++ b/src/Pipelines/ZeroShotImageClassificationPipeline.php
@@ -5,7 +5,6 @@
 
 namespace Codewithkyrian\Transformers\Pipelines;
 
-use Codewithkyrian\Transformers\Utils\Math;
 use Codewithkyrian\Transformers\Utils\Tensor;
 use function Codewithkyrian\Transformers\Utils\prepareImages;
 
@@ -53,14 +52,14 @@ public function __invoke(array|string $inputs, ...$args): array
 
         $activationFn = $this->model->config['model_type'] === 'siglip' ?
             fn(Tensor $batch) => $batch->sigmoid()->toArray() :
-            fn(Tensor $batch) => Math::softmax($batch->toArray());
+            fn(Tensor $batch) => $batch->softmax();
 
         // Compare each image with each candidate label
         $toReturn = [];
 
         foreach ($output['logits_per_image'] as $batch) {
             // Compute softmax per image
-            $scores = $activationFn(Tensor::fromArray($batch));
+            $scores = $activationFn($batch);
 
             $result = [];
             foreach ($scores as $i => $score) {
diff --git a/src/Processors/Processor.php b/src/Processors/Processor.php
index c9f45e0..800f537 100644
--- a/src/Processors/Processor.php
+++ b/src/Processors/Processor.php
@@ -8,7 +8,6 @@
 use Codewithkyrian\Transformers\FeatureExtractors\FeatureExtractor;
 use Codewithkyrian\Transformers\Models\Output\ObjectDetectionOutput;
 use Codewithkyrian\Transformers\Utils\Math;
-use Codewithkyrian\Transformers\Utils\Tensor;
 use Exception;
 
 /**
@@ -70,10 +69,8 @@ public static function postProcessObjectDetection(ObjectDetectionOutput $outputs
                     }
                     $probs = $logitSigmoid;
                 } else {
-                    $mo = Tensor::getMo();
-
                     // Get most probable class
-                    $maxIndex = $mo->argMax($logit);
+                    $maxIndex = $logit->argMax();
 
                     if ($maxIndex === $numClasses - 1) {
                         // This is the background class, skip it
@@ -82,7 +79,7 @@ public static function postProcessObjectDetection(ObjectDetectionOutput $outputs
                     $indices[] = $maxIndex;
 
                     // Compute softmax over classes
-                    $probs = Math::softmax($logit->toArray());
+                    $probs = $logit->softmax();
                 }
 
                 foreach ($indices as $index) {
diff --git a/src/Utils/Helpers.php b/src/Utils/Helpers.php
index f5227e4..85cc2d6 100644
--- a/src/Utils/Helpers.php
+++ b/src/Utils/Helpers.php
@@ -4,8 +4,6 @@
 
 namespace Codewithkyrian\Transformers\Utils;
 
-use Exception;
-
 function memoryUsage(): string
 {
     $mem = memory_get_usage(true);
diff --git a/src/Utils/Tensor.php b/src/Utils/Tensor.php
index 0a31b46..661e497 100644
--- a/src/Utils/Tensor.php
+++ b/src/Utils/Tensor.php
@@ -929,9 +929,22 @@ protected function unflattenArray($flatArray, &$currentIndex, array $shape): arr
      * Calculate the softmax of the tensor.
      *
      */
-    public function softmax(): array
+    public function softmax(): array|static
     {
-        return Math::softmax($this->toArray());
+        return match ($this->ndim()) {
+            1 => Math::softmax($this->toArray()),
+            2 => $this->softmax2D(),
+            default => throw new InvalidArgumentException("Softmax is only supported for 1D and 2D tensors.")
+        };
+    }
+
+    protected function softmax2D(): static
+    {
+        $mo = self::getMo();
+
+        $ndArray = $mo->la()->softmax($this);
+
+        return new static($ndArray->buffer(), $ndArray->dtype(), $ndArray->shape(), $ndArray->offset());
     }
 
     public function max(?int $axis = null): static|int|float