Adding power and log softmax operations (#80)

* Adding power and log softmax operations

* Add functionality to handle tensor params for power opeation

* Fix of function comments

---------

Co-authored-by: SarahByrneIntel <[email protected]>

Author: SarahByrneIntel

include/intel_npu_acceleration_library/nn_factory.h

@@ -884,6 +884,33 @@ class ModelFactory : public intel_npu_acceleration_library::OVInferenceModel {
         return normL2.get();
     }
 
+    /**
+     * @brief Create a new power operation
+     *
+     * @param x1 operation's input node
+     * @param x2 operation's input node of the exponent
+     * @param auto_broadcast auto broadcast specification
+     * @return ov::op::Op*
+     */
+    ov::op::Op* power(ov::op::Op* x1, ov::op::Op* x2, ov::op::AutoBroadcastType auto_broadcast) {
+        auto power = std::make_shared<ov::opset1::Power>(x1->output(0), x2->output(0), auto_broadcast);
+        operations.push_back(power);
+        return power.get();
+    }
+
+    /**
+     * @brief Create a new log softmax operation
+     *
+     * @param input operation's input node
+     * @param axis the axis position on which to calculate the LogSoftmax
+     * @return ov::op::Op*
+     */
+    ov::op::Op* log_softmax(ov::op::Op* input, int64_t axis) {
+        auto log_softmax = std::make_shared<ov::opset5::LogSoftmax>(input->output(0), axis);
+        operations.push_back(log_softmax);
+        return log_softmax.get();
+    }
+
     void result(ov::op::Op* op) {
         auto res = std::make_shared<ov::opset8::Result>(op->output(0));
         results.push_back(res);

intel_npu_acceleration_library/backend/factory.py

@@ -563,6 +563,41 @@ def normL2(
         axis_node = self.constant(axis).node  # type: ignore
         return backend_lib.normL2(self._mm, input_node, axis_node, eps)
 
+    @return_tensor
+    def power(
+        self,
+        input_node: ctypes._Pointer,
+        exponent: Union[ctypes._Pointer, torch.Tensor],
+    ) -> ctypes._Pointer:
+        """Generate a power layer.
+
+        Args:
+            input_node (ctypes._Pointer): layer input node
+            exponent (Union[ctypes._Pointer, torch.Tensor]): the exponent value
+
+        Raises:
+            ValueError: Input tensor shapes are not equal
+
+        Returns:
+            ctypes._Pointer: output node
+        """
+        input_shape_size = backend_lib.op_shape_size(input_node)
+        input_shape = [
+            backend_lib.op_shape(input_node, i) for i in range(input_shape_size)
+        ]
+        if isinstance(exponent, ctypes._Pointer):
+            exponent_shape_size = backend_lib.op_shape_size(input_node)
+            exponent_shape = [
+                backend_lib.op_shape(exponent, i) for i in range(exponent_shape_size)
+            ]
+        else:
+            exponent_shape = list(exponent.shape)
+            exponent = self.constant(exponent).node  # type: ignore
+        if exponent_shape != input_shape:
+            raise ValueError("Input tensor shapes are not equal")
+
+        return backend_lib.power(self._mm, input_node, exponent)
+
     @return_tensor
     def avg_pooling(
         self,

intel_npu_acceleration_library/backend/ops.py

@@ -132,5 +132,7 @@ def get_supported_ops() -> List[SupportedOp]:
         ),
         SupportedOp(name="adaptive_avg_pool", inputs=2),
         SupportedOp(name="adaptive_max_pool", inputs=2),
+        SupportedOp(name="power", inputs=2),
+        SupportedOp(name="log_softmax", inputs=1, parameters=[ctypes.c_int64]),
     ]
     return supported_ops

intel_npu_acceleration_library/nn/functional.py

@@ -1109,3 +1109,46 @@ def conv2d(
     )
 
     return conv
+
+
+@implements(torch.pow)
+def pow(input: Tensor, exponent: Union[Tensor, torch.Tensor, float]):
+    """Return the tensor raised to the power of the exponent.
+
+    Args:
+        input (Tensor): The input tensor.
+        exponent (Union[Tensor, torch.Tensor, float]): The exponent value.
+
+    Returns:
+        Tensor: Output tensor.
+    """
+    if isinstance(exponent, float):
+        exponent = torch.full(input.shape, exponent).to(torch.float16)
+    return generate_op([input], "power", exponent=exponent)
+
+
+@implements(torch.nn.functional.log_softmax)
+def log_softmax(
+    input: Tensor,
+    dim: Optional[int] = None,
+    _stacklevel=3,
+    dtype: Optional[torch.dtype] = None,
+) -> Tensor:
+    """Return the log softmax of a tensor element-wise.
+
+    Args:
+        input (Tensor): The input tensor.
+        dim (int): The dimension along which log_softmax will be computed. Defaults to -1.
+        _stacklevel (int): The stack level. Defaults to 3.
+        dtype (torch.dtype): The data type. Defaults to None.
+
+    Returns:
+        Tensor: Output tensor.
+    """
+    if dim is None:
+        dim = -1
+    log_smax = generate_op([input], "log_softmax", dim)
+
+    if dtype:
+        log_smax = log_smax.to(dtype)
+    return log_smax

src/bindings.cpp

@@ -563,4 +563,14 @@ intel_npu_acceleration_library_DLL_API ov::op::Op* normL2(intel_npu_acceleration
                                                           ov::op::Op* data, ov::op::Op* axes, float eps) {
     return factory->normL2(data, axes, eps);
 }
+
+intel_npu_acceleration_library_DLL_API ov::op::Op* power(intel_npu_acceleration_library::ModelFactory* factory,
+                                                         ov::op::Op* x1, ov::op::Op* x2) {
+    return factory->power(x1, x2, ov::op::AutoBroadcastType::NUMPY);
+}
+
+intel_npu_acceleration_library_DLL_API ov::op::Op* log_softmax(intel_npu_acceleration_library::ModelFactory* factory,
+                                                               ov::op::Op* input, int64_t axis) {
+    return factory->log_softmax(input, axis);
+}
 }

test/python/test_op.py

@@ -452,3 +452,58 @@ def test_multiple_outputs():
 
     assert 1 - r2_score(result0.detach().numpy().flatten(), ref1.flatten()) < 0.01
     assert 1 - r2_score(result1.detach().numpy().flatten(), ref2.flatten()) < 0.01
+
+
+@pytest.mark.parametrize("batch", [16, 128])
+@pytest.mark.parametrize("hidden_dim", [128, 256])
+@pytest.mark.parametrize("exponent", ["tensor", "float"])
+@pytest.mark.parametrize("exponent_type", ["parameter", "constant"])
+def test_power(batch, hidden_dim, exponent, exponent_type):
+
+    x = torch.rand((batch, hidden_dim)).to(torch.float16)
+    if exponent == "tensor":
+        exponent = torch.rand((batch, hidden_dim)).to(torch.float16)
+    else:
+        exponent = torch.rand(1).to(torch.float16).item()
+
+    reference = torch.pow(x, exponent=exponent).numpy()
+
+    model = NNFactory()
+    par = model.parameter(x.shape, np.float16)
+    if exponent == "tensor" and exponent_type == "parameter":
+        exponent_par = model.parameter(exponent.shape, np.float16)
+        _ = torch.pow(par, exponent_par)
+        model.compile()
+        out = model(x, exponent).numpy()
+    else:
+        _ = torch.pow(par, exponent=exponent)
+        model.compile()
+        out = model(x).numpy()
+
+    assert out.shape == reference.shape, "Output shape mismatch"
+    assert np.isfinite(reference).all(), "Pytorch Reference contains NaN or Inf"
+    assert np.isfinite(out).all(), "NPU output contains NaN or Inf"
+
+    assert 1 - r2_score(reference, out) < 0.01
+
+
+@pytest.mark.parametrize("batch", [16, 128])
+@pytest.mark.parametrize("hidden_dim", [128, 256])
+@pytest.mark.parametrize("axis", [0, 1, -1, -2])
+def test_logsoftmax(batch, hidden_dim, axis):
+    x = torch.rand((batch, hidden_dim)).to(torch.float16)
+
+    reference = torch.nn.functional.log_softmax(x, dim=axis).numpy()
+
+    model = NNFactory()
+    par = model.parameter(x.shape, np.float16)
+    _ = torch.nn.functional.log_softmax(par, dim=axis)
+    model.compile()
+
+    out = model(x).numpy()
+
+    assert out.shape == reference.shape, "Output shape mismatch"
+    assert np.isfinite(reference).all(), "Pytorch Reference contains NaN or Inf"
+    assert np.isfinite(out).all(), "NPU output contains NaN or Inf"
+
+    assert 1 - r2_score(reference, out) < 0.01