From 7ec4128d33ee19907ac9b9bdd3a1043beb225c9a Mon Sep 17 00:00:00 2001
From: usmanmateen <74776275+usmanmateen@users.noreply.github.com>
Date: Sun, 11 Oct 2026 01:38:48 +0100
Subject: [PATCH] Support double in optimized log_softmax

The optimized _log_softmax.out kernel rejected float64 inputs with
InvalidArgument, while the portable kernel supports them. The ATen
vectorized helpers it uses work for double, so run it through the same
same-dtype path as Half and BFloat16, and stop skipping the double tests
for the optimized kernels.
---
 kernels/optimized/cpu/op_log_softmax.cpp           | 8 ++++++--
 kernels/optimized/test/supported_features_def.yaml | 3 ---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/kernels/optimized/cpu/op_log_softmax.cpp b/kernels/optimized/cpu/op_log_softmax.cpp
index d49e72efcc2..a9117216ac6 100644
--- a/kernels/optimized/cpu/op_log_softmax.cpp
+++ b/kernels/optimized/cpu/op_log_softmax.cpp
@@ -102,6 +102,7 @@ void log_softmax_kernel(const Tensor& input, int64_t dim, Tensor& out) {
 template <typename OUT_T>
 bool log_softmax_wrapper(const Tensor& X, int64_t dim, Tensor& out) {
   if constexpr (
+      std::is_same_v<OUT_T, double> ||
       std::is_same_v<OUT_T, executorch::aten::BFloat16> ||
       std::is_same_v<OUT_T, executorch::aten::Half>) {
     // Input dtype equals output dtype (enforced by check_log_softmax_args).
@@ -112,7 +113,6 @@ bool log_softmax_wrapper(const Tensor& X, int64_t dim, Tensor& out) {
   } else {
     auto input_scalar_type = X.scalar_type();
     switch (input_scalar_type) {
-      // TODO: support Double as well
       case ScalarType::Float:
         log_softmax_kernel<float, OUT_T>(X, dim, out);
         return true;
@@ -149,12 +149,16 @@ Tensor& opt_log_softmax_out(
 
   auto out_scalar_type = out.scalar_type();
   switch (out_scalar_type) {
-    // TODO: support Double as well
     case ScalarType::Float: {
       bool success = log_softmax_wrapper<float>(self, dim, out);
       ET_KERNEL_CHECK(context, success, InvalidArgument, out);
       break;
     }
+    case ScalarType::Double: {
+      bool success = log_softmax_wrapper<double>(self, dim, out);
+      ET_KERNEL_CHECK(context, success, InvalidArgument, out);
+      break;
+    }
     case ScalarType::BFloat16: {
       bool success =
           log_softmax_wrapper<executorch::aten::BFloat16>(self, dim, out);
diff --git a/kernels/optimized/test/supported_features_def.yaml b/kernels/optimized/test/supported_features_def.yaml
index df51f3ec2f4..054ba32e29a 100644
--- a/kernels/optimized/test/supported_features_def.yaml
+++ b/kernels/optimized/test/supported_features_def.yaml
@@ -1,5 +1,2 @@
 - namespace: op_gelu
   dtype_double: false
-
-- namespace: op_log_softmax
-  dtype_double: false