import torch
import ksc.torch_frontend as knossos

# run-bench: Knossos source, and "nice" PyTorch implementation
# BEGINDOC
@knossos.register
def sqrl(x: torch.Tensor):
    """
    sqrl: Squared Leaky Relu
    Like a capsule from /Stuck in a Rut/
    Typically x is a 4x4 tensor, possibly
    packed in a 4n x 4m array
    """
    y = torch.sum(x)
    if y < 0.0:
        t = -0.125 * x
    else:
        t = 1 / 2 * x ** 2
    return torch.mean(torch.sin(t) * t)


# ENDDOC

# run-bench: PyTorch "fast" implementation
def sqrl_pytorch(x: torch.Tensor):
    return sqrl(x)


# run-bench: PyTorch "nice" implementation
def sqrl_pytorch_nice(x: torch.Tensor):
    return sqrl(x)


# run-bench: Define a range of values at which to call the methods
def sqrl_bench_configs():
    yield torch.randn((4, 4))
    yield torch.randn((16, 16))


#################################
#
# vsqrl - vectorized sqrl
#

vsqrl = knossos.vmap(sqrl)


# run-bench: Define a range of values at which to call the methods
def vsqrl_bench_configs():
    yield torch.randn((10, 4, 4))
    yield torch.randn((1000, 4, 4))
    yield torch.randn((1000, 16, 16))

examples/dl-capsule/sqrl.py
import torch
import ksc.torch_frontend as knossos

# run-bench: Knossos source, and "nice" PyTorch implementation
# BEGINDOC
@knossos.register
def sqrl(x: torch.Tensor):
    """
    sqrl: Squared Leaky Relu
    Like a capsule from /Stuck in a Rut/
    Typically x is a 4x4 tensor, possibly
    packed in a 4n x 4m array
    """
    y = torch.sum(x)
    if y < 0.0:
        t = -0.125 * x
    else:
        t = 1 / 2 * x ** 2
    return torch.mean(torch.sin(t) * t)


# ENDDOC

# run-bench: PyTorch "fast" implementation
def sqrl_pytorch(x: torch.Tensor):
    return sqrl.raw_f(x)


# run-bench: PyTorch "nice" implementation
def sqrl_pytorch_nice(x: torch.Tensor):
    return sqrl.raw_f(x)


# run-bench: Define a range of values at which to call the methods
def sqrl_bench_configs():
    yield torch.randn((4, 4))
    yield torch.randn((16, 16))


#################################
#
# vsqrl - vectorized sqrl
#

vsqrl = knossos.vmap(sqrl)


# run-bench: Define a range of values at which to call the methods
def vsqrl_bench_configs():
    yield torch.randn((10, 4, 4))
    yield torch.randn((1000, 4, 4))
    yield torch.randn((1000, 16, 16))

examples/dl-capsule/sqrl.py
diff --git a/examples/dl-capsule/sqrl.py b/examples/dl-capsule/sqrl.py
index f4f94791e..2f022b265 100644
--- a/examples/dl-capsule/sqrl.py
+++ b/examples/dl-capsule/sqrl.py
@@ -23,12 +23,12 @@ def sqrl(x: torch.Tensor):
 
 # run-bench: PyTorch "fast" implementation
 def sqrl_pytorch(x: torch.Tensor):
-    return sqrl(x)
+    return sqrl.raw_f(x)
 
 
 # run-bench: PyTorch "nice" implementation
 def sqrl_pytorch_nice(x: torch.Tensor):
-    return sqrl(x)
+    return sqrl.raw_f(x)
 
 
 # run-bench: Define a range of values at which to call the methods

Bug: Segmentation fault in sqrl_pytorch-PyTorch CUDA
Just saw this while working on something else.   I haven't done a lot to debug it, but note that it's in copydown, on a fairly innocuous operation (aten::sum(Tensor 2) -> Float), so might be something to do with KS_ALLOCATOR not being defined?
Or could just be out of memory not caught?
![image](https://user-images.githubusercontent.com/128119/130364073-023976aa-9222-4d55-ad43-95a7ac530d76.png)

