Add float_power operator

Signed-off-by: xuanzhengdu-eng <xuanzhengdu@gmail.com>

Author: xuanzhengdu-eng

benchmark/test_float_power.py

@@ -0,0 +1,122 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from . import base, consts
+
+
+class FloatPowerBenchmark(base.GenericBenchmark):
+    def set_shapes(self, shape_file_path=None):
+        self.shapes = [(4,), (1024,), (65536,), (512, 512), (2048, 2048)]
+
+
+def _tensor_tensor_input(shape, dtype, device):
+    base_tensor = torch.rand(shape, dtype=dtype, device=device).add_(0.25)
+    exponent = torch.rand(shape, dtype=dtype, device=device).mul_(4).sub_(2)
+    yield base_tensor, exponent
+
+
+def _tensor_scalar_input(shape, dtype, device):
+    base_tensor = torch.rand(shape, dtype=dtype, device=device).add_(0.25)
+    yield base_tensor, 1.234
+
+
+def _scalar_tensor_input(shape, dtype, device):
+    exponent = torch.rand(shape, dtype=dtype, device=device).mul_(4).sub_(2)
+    yield 2.0, exponent
+
+
+def _tensor_tensor_out_input(shape, dtype, device):
+    base_tensor = torch.rand(shape, dtype=dtype, device=device).add_(0.25)
+    exponent = torch.rand(shape, dtype=dtype, device=device).mul_(4).sub_(2)
+    out = torch.empty(shape, dtype=torch.float64, device=device)
+    yield base_tensor, exponent, {"out": out}
+
+
+def _tensor_scalar_out_input(shape, dtype, device):
+    base_tensor = torch.rand(shape, dtype=dtype, device=device).add_(0.25)
+    out = torch.empty(shape, dtype=torch.float64, device=device)
+    yield base_tensor, 1.234, {"out": out}
+
+
+def _scalar_tensor_out_input(shape, dtype, device):
+    exponent = torch.rand(shape, dtype=dtype, device=device).mul_(4).sub_(2)
+    out = torch.empty(shape, dtype=torch.float64, device=device)
+    yield 2.0, exponent, {"out": out}
+
+
+def _run_benchmark(op_name, input_fn, torch_op):
+    bench = FloatPowerBenchmark(
+        input_fn=input_fn,
+        op_name=op_name,
+        torch_op=torch_op,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.float_power_tensor_tensor
+def test_float_power_tensor_tensor():
+    _run_benchmark(
+        "float_power_tensor_tensor",
+        _tensor_tensor_input,
+        torch.ops.aten.float_power.Tensor_Tensor,
+    )
+
+
+@pytest.mark.float_power_tensor_scalar
+def test_float_power_tensor_scalar():
+    _run_benchmark(
+        "float_power_tensor_scalar",
+        _tensor_scalar_input,
+        torch.ops.aten.float_power.Tensor_Scalar,
+    )
+
+
+@pytest.mark.float_power_scalar_tensor
+def test_float_power_scalar_tensor():
+    _run_benchmark(
+        "float_power_scalar_tensor",
+        _scalar_tensor_input,
+        torch.ops.aten.float_power.Scalar,
+    )
+
+
+@pytest.mark.float_power_tensor_tensor_out
+def test_float_power_tensor_tensor_out():
+    _run_benchmark(
+        "float_power_tensor_tensor_out",
+        _tensor_tensor_out_input,
+        torch.ops.aten.float_power.Tensor_Tensor_out,
+    )
+
+
+@pytest.mark.float_power_tensor_scalar_out
+def test_float_power_tensor_scalar_out():
+    _run_benchmark(
+        "float_power_tensor_scalar_out",
+        _tensor_scalar_out_input,
+        torch.ops.aten.float_power.Tensor_Scalar_out,
+    )
+
+
+@pytest.mark.float_power_scalar_tensor_out
+def test_float_power_scalar_tensor_out():
+    _run_benchmark(
+        "float_power_scalar_tensor_out",
+        _scalar_tensor_out_input,
+        torch.ops.aten.float_power.Scalar_out,
+    )

conf/operators.yaml

@@ -3542,6 +3542,72 @@ ops:
       - Math
     stages:
       - alpha: '5.4'
+  - id: float_power_scalar_tensor
+    description: Computes a scalar base raised element-wise to tensor exponents in float64.
+    for:
+      - float_power.Scalar
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - Math
+    stages:
+      - alpha: '5.4'
+  - id: float_power_scalar_tensor_out
+    description: Out variant of float_power.Scalar.
+    for:
+      - float_power.Scalar_out
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - Math
+    stages:
+      - alpha: '5.4'
+  - id: float_power_tensor_scalar
+    description: Raises tensor elements to a scalar exponent in float64.
+    for:
+      - float_power.Tensor_Scalar
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - Math
+    stages:
+      - alpha: '5.4'
+  - id: float_power_tensor_scalar_out
+    description: Out variant of float_power.Tensor_Scalar.
+    for:
+      - float_power.Tensor_Scalar_out
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - Math
+    stages:
+      - alpha: '5.4'
+  - id: float_power_tensor_tensor
+    description: Raises tensor elements to tensor exponents in float64 with broadcasting.
+    for:
+      - float_power.Tensor_Tensor
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - Math
+    stages:
+      - alpha: '5.4'
+  - id: float_power_tensor_tensor_out
+    description: Out variant of float_power.Tensor_Tensor.
+    for:
+      - float_power.Tensor_Tensor_out
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - Math
+    stages:
+      - alpha: '5.4'
   - id: floor
     description: |
       Performs an element-wise floor operation, rounding each element of a tensor

src/flag_gems/__init__.py

@@ -478,6 +478,12 @@ def torch_ge(v):
     ("fix_", fix_),
     ("flatten.using_ints", flatten),
     ("flip", flip),
+    ("float_power.Scalar", float_power_scalar_tensor),
+    ("float_power.Scalar_out", float_power_scalar_tensor_out),
+    ("float_power.Tensor_Scalar", float_power_tensor_scalar),
+    ("float_power.Tensor_Scalar_out", float_power_tensor_scalar_out),
+    ("float_power.Tensor_Tensor", float_power_tensor_tensor),
+    ("float_power.Tensor_Tensor_out", float_power_tensor_tensor_out),
     ("float_power_.Scalar", float_power_tensor_scalar_),
     ("float_power_.Tensor", float_power_tensor_tensor_),
     ("floor", floor),

src/flag_gems/ops/__init__.py

@@ -303,10 +303,14 @@
 from flag_gems.ops.flatten import flatten
 from flag_gems.ops.flip import flip
 from flag_gems.ops.float_power_ import (
+    float_power_scalar_tensor,
+    float_power_scalar_tensor_out,
     float_power_tensor_scalar,
     float_power_tensor_scalar_,
+    float_power_tensor_scalar_out,
     float_power_tensor_tensor,
     float_power_tensor_tensor_,
+    float_power_tensor_tensor_out,
 )
 from flag_gems.ops.floor import floor, floor_out
 from flag_gems.ops.floor_ import floor_
@@ -1117,10 +1121,14 @@
     "flash_attn_varlen_func",
     "flash_attn_varlen_opt_func",
     "flip",
+    "float_power_scalar_tensor",
+    "float_power_scalar_tensor_out",
     "float_power_tensor_scalar",
     "float_power_tensor_scalar_",
+    "float_power_tensor_scalar_out",
     "float_power_tensor_tensor",
     "float_power_tensor_tensor_",
+    "float_power_tensor_tensor_out",
     "floor",
     "floor_",
     "floor_divide",

src/flag_gems/ops/float_power_.py

@@ -1,6 +1,7 @@
 # Generated by KernelGen: https://github.com/flagos-ai/KernelGen
 import logging
 
+import torch
 import triton
 import triton.language as tl
 
@@ -18,7 +19,12 @@ def float_power_func(x, exponent):
 
 def float_power_tensor_tensor(A, exponent):
     logger.debug("GEMS FLOAT_POWER_TENSOR_TENSOR")
-    return float_power_func(A, exponent)
+    out = torch.empty(
+        torch.broadcast_shapes(A.shape, exponent.shape),
+        dtype=torch.float64,
+        device=A.device,
+    )
+    return float_power_tensor_tensor_func(A, exponent, out0=out)
 
 
 def float_power_tensor_tensor_(A, exponent):
@@ -34,9 +40,68 @@ def float_power_func_tensor_scalar(x, exponent):
 
 def float_power_tensor_scalar(A, exponent):
     logger.debug("GEMS FLOAT_POWER_TENSOR_SCALAR")
-    return float_power_func_tensor_scalar(A, exponent)
+    out = torch.empty(A.shape, dtype=torch.float64, device=A.device)
+    return float_power_tensor_scalar_func(A, exponent, out0=out)
 
 
 def float_power_tensor_scalar_(A, exponent):
     logger.debug("GEMS FLOAT_POWER_TENSOR_SCALAR_")
     return float_power_func_tensor_scalar(A, exponent, out0=A)
+
+
+@pointwise_dynamic(promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def float_power_tensor_tensor_func(x, exponent):
+    return _pow(x.to(tl.float64), exponent.to(tl.float64))
+
+
+@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def float_power_tensor_scalar_func(x, exponent):
+    return _pow(x.to(tl.float64), exponent.to(tl.float64))
+
+
+@pointwise_dynamic(is_tensor=[False, True], promotion_methods=[(0, 1, "DEFAULT")])
+@triton.jit
+def float_power_scalar_tensor_func(x, exponent):
+    return _pow(x.to(tl.float64), exponent.to(tl.float64))
+
+
+def _prepare_out(out, shape, device):
+    if out.dtype != torch.float64:
+        raise RuntimeError(
+            f"the output given to float_power has dtype {out.dtype} "
+            "but the operation's result requires dtype Double"
+        )
+    if out.device != device:
+        raise RuntimeError(
+            f"Expected out tensor to have device {device}, but got {out.device} instead"
+        )
+    if out.shape != shape:
+        out.resize_(shape)
+    return out
+
+
+def float_power_tensor_tensor_out(A, exponent, *, out):
+    logger.debug("GEMS FLOAT_POWER_TENSOR_TENSOR_OUT")
+    shape = torch.broadcast_shapes(A.shape, exponent.shape)
+    _prepare_out(out, shape, A.device)
+    return float_power_tensor_tensor_func(A, exponent, out0=out)
+
+
+def float_power_tensor_scalar_out(A, exponent, *, out):
+    logger.debug("GEMS FLOAT_POWER_TENSOR_SCALAR_OUT")
+    _prepare_out(out, A.shape, A.device)
+    return float_power_tensor_scalar_func(A, exponent, out0=out)
+
+
+def float_power_scalar_tensor(A, exponent):
+    logger.debug("GEMS FLOAT_POWER_SCALAR_TENSOR")
+    out = torch.empty(exponent.shape, dtype=torch.float64, device=exponent.device)
+    return float_power_scalar_tensor_func(A, exponent, out0=out)
+
+
+def float_power_scalar_tensor_out(A, exponent, *, out):
+    logger.debug("GEMS FLOAT_POWER_SCALAR_TENSOR_OUT")
+    _prepare_out(out, exponent.shape, exponent.device)
+    return float_power_scalar_tensor_func(A, exponent, out0=out)