nll libtorch

Author: ACSadighi

bridge/include/bridge.h

@@ -72,6 +72,14 @@ bridge_tensor_t conv2d(
     int padding
 );
 
+bridge_tensor_t nll_loss(
+    bridge_tensor_t input, 
+    bridge_tensor_t target, 
+    bridge_tensor_t weight,
+    int ignoreIndex,
+    int reduction
+);
+
 bridge_tensor_t matmul(bridge_tensor_t a, bridge_tensor_t b);
 
 bridge_tensor_t max_pool2d(

bridge/lib/bridge.cpp

@@ -279,6 +279,39 @@ extern "C" bridge_tensor_t add_two_arrays(bridge_tensor_t a, bridge_tensor_t b)
     return torch_to_bridge(output);
 }
 
+extern "C" bridge_tensor_t nll_loss(
+    bridge_tensor_t input, 
+    bridge_tensor_t target, 
+    bridge_tensor_t weight,
+    int ignoreIndex,
+    int reduction
+) {
+    // Convert bridge_tensor_t to torch::Tensor
+    torch::Tensor t_input = bridge_to_torch(input);
+    torch::Tensor t_target = bridge_to_torch(target);
+    torch::Tensor t_weight = bridge_to_torch(weight);
+
+    // Map reduction int to string
+    std::string reduction_str;
+    switch (reduction) {
+        case 0: reduction_str = "none"; break;
+        case 1: reduction_str = "mean"; break;
+        case 2: reduction_str = "sum"; break;
+        default: reduction_str = "mean"; break;  // fallback default
+    }
+
+    torch::Tensor output = torch::nn::functional::nll_loss(
+        t_input,
+        t_target,
+        torch::nn::functional::NLLLossFuncOptions()
+            .weight(t_weight)
+            .ignore_index(ignoreIndex)
+            .reduction(reduction_str)
+    );
+
+    return torch_to_bridge(output);
+}
+
 // extern "C" bridge_tensor_t capture_webcam_bridge(int cam_index) {
 //     torch::Tensor image = capture_webcam(cam_index);
 //     return torch_to_bridge(image);

lib/Autograd.chpl

@@ -1084,13 +1084,12 @@ record nllLossOp : serializable {
     var target: shared BaseTensorResource(?);
     var weight: shared BaseTensorResource(?);
     var ignoreIndex: int;
-    var red: bool;
     var reduction: string;
 
     proc children do return (input,target,weight);
 
     proc forward() do
-        return ndarray.nllLoss(input.array,target.array,weight.array,ignoreIndex,red,reduction);
+        return ndarray.nllLoss(input.array,target.array,weight.array,ignoreIndex,reduction);
 
     proc spec : GradOpSpec do return new dict(("operation", "nllLoss"));
 }

lib/Bridge.chpl

@@ -87,6 +87,13 @@ module Bridge {
         in a: bridge_tensor_t, 
         in b: bridge_tensor_t): bridge_tensor_t;
 
+    extern "nll_loss" proc nllLoss(
+        in input: bridge_tensor_t, 
+        in target: ndarray(1,eltType), 
+        in weight: ndarray(1, eltType),
+        in ignoreIndex: int(32),
+        in reduction: int(32): bridge_tensor_t;
+
     extern "split_loop" proc splitLoop(idx: int(64), n: int(64)): void;
 
     extern "split_loop_filler" proc splitLoopFiller(n: int(64),ret: c_ptr(int(64))): void;

lib/DynamicTensor.chpl

@@ -733,14 +733,13 @@ proc type dynamicTensor.nllLoss(
     target: dynamicTensor(eltType), 
     weight: dynamicTensor(eltType),
     ignoreIndex: int = -1,
-    red: bool = true,
     reduction: string = "mean"
 ) {
     for param rankIn in 2..2 {
         if input.checkRank(rankIn) {
             for param rank in 1..1 {
                 if target.checkRank(rankIn) && weight.checkRank(rank) {
-                    return staticTensor.nllLoss(input.forceRank(rankIn),target.forceRank(rank),weight.forceRank(rank),ignoreIndex,red,reduction);
+                    return staticTensor.nllLoss(input.forceRank(rankIn),target.forceRank(rank),weight.forceRank(rank),ignoreIndex,reduction);
                 }
             }
         }
@@ -751,7 +750,6 @@ proc type dynamicTensor.nllLoss(
     input: dynamicTensor(?eltType), 
     target: dynamicTensor(eltType), 
     ignoreIndex: int = -1,
-    red: bool = true,
     reduction: string = "mean"
 ) {
     param inRank: int = 2;
@@ -762,7 +760,7 @@ proc type dynamicTensor.nllLoss(
             var stInput: staticTensor(inRank,eltType) = input.forceRank(inRank);
             var stTarget: staticTensor(targetRank,eltType) = target.forceRank(targetRank);
             var weights: staticTensor(1,eltType) = staticTensor.ones(eltType,3);
-            return staticTensor.nllLoss(stInput,stTarget,weights,ignoreIndex,red,reduction);
+            return staticTensor.nllLoss(stInput,stTarget,weights,ignoreIndex,reduction);
         }
     }
 

lib/NDArray.chpl

@@ -2140,36 +2140,46 @@ proc type ndarray.nllLoss(
     target: ndarray(1,eltType), 
     weight: ndarray(1, eltType),
     ignoreIndex: int = -1,
-    red: bool = true,
     reduction: string = "mean"
 ): ndarray(1,eltType) {
-    const (N,C) = input.shape;
-    assert(target.shape[0] == N, "Target shape must match batch size.");
-    assert(weight.shape[0] == C, "Weights shape must match number of classes.");
+    int reduction_int = 1;
+    if reduction == "sum" then reduction_int = 2;
+    if reduction == "none" then reduction_int = 0;
+
+    return Bridge.nllLoss(
+        input: Bridge.tensorHandle(eltType), 
+        target: Bridge.tensorHandle(eltType), 
+        weight: Bridge.tensorHandle(eltType),
+        ignoreIndex,
+        reduction
+    ) : ndarray(rank,eltType);
+    // const (N,C) = input.shape;
+    // assert(target.shape[0] == N, "Target shape must match batch size.");
+    // assert(weight.shape[0] == C, "Weights shape must match number of classes.");
 
-    const dom = util.domainFromShape(N);
-    var loss = new ndarray(dom, eltType);
-    ref x = input.data;
-    ref y = target.data;
-    ref w = weight.data;
-    ref lossD = loss.data;
-    var wynSum: real = 0.0;
-
-    forall n in 0..<N with (+ reduce wynSum) {
-        const yn: int = y[n]:int;
-        if yn == ignoreIndex {
-            lossD[n] = 0.0;
-        }
-        else {
-            lossD[n] = -w[yn]*x[n,yn];
-            wynSum += w[yn];
-        }
-    }
+    // const dom = util.domainFromShape(N);
+    // var loss = new ndarray(dom, eltType);
+    // ref x = input.data;
+    // ref y = target.data;
+    // ref w = weight.data;
+    // ref lossD = loss.data;
+    // var wynSum: real = 0.0;
+
+    // forall n in 0..<N with (+ reduce wynSum) {
+    //     const yn: int = y[n]:int;
+    //     if yn == ignoreIndex {
+    //         lossD[n] = 0.0;
+    //     }
+    //     else {
+    //         lossD[n] = -w[yn]*x[n,yn];
+    //         wynSum += w[yn];
+    //     }
+    // }
 
-    if !red then return loss;
-    if reduction == "mean" then return loss.sum(0) / wynSum;
-    if reduction == "sum" then return loss.sum(0);
-    halt("Invalid reduction mode: " + reduction);
+    // if !red then return loss;
+    // if reduction == "mean" then return loss.sum(0) / wynSum;
+    // if reduction == "sum" then return loss.sum(0);
+    // halt("Invalid reduction mode: " + reduction);
 }
 
 module ndarrayRandom {

lib/StaticTensor.chpl

@@ -517,10 +517,9 @@ proc type staticTensor.nllLoss(
     target: staticTensor(1,eltType),
     weight: staticTensor(1,eltType),
     ignoreIndex: int = -1,
-    red: bool = true,
     reduction: string = "mean"
 ) {
-    var ctx = new nllLossOp(input.meta,target.meta,weight.meta,ignoreIndex,red,reduction);
+    var ctx = new nllLossOp(input.meta,target.meta,weight.meta,ignoreIndex,reduction);
     return tensorFromCtx(1,eltType,ctx);
 }