From 8284fd9f73e553e3a36c45d69cea2375cf91aec0 Mon Sep 17 00:00:00 2001 From: ssjia Date: Tue, 10 Mar 2026 10:01:20 -0700 Subject: [PATCH] [ET-VK] Add ANY_STORAGE support to argmax and argmin Register argmax and argmin as ANY_STORAGE. The existing shaders and C++ dispatch already handle buffer storage via pick_storage_for_reduce. Differential Revision: [D95970165](https://our.internmc.facebook.com/intern/diff/D95970165/) [ghstack-poisoned] --- backends/vulkan/op_registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 056869ecc98..13bd4acbfef 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -703,7 +703,7 @@ def register_reduce_cpp_ops(): ) def register_argreduce_cpp_ops(): return OpFeatures( - inputs_storage=utils.ANY_TEXTURE, + inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_T, outputs_dtypes=utils.INT_T, supports_resize=True,