Package: slope-ad
SlopeSlope is a small automatic differentation (AD) engine, focused on machine learning (ML)
Tensor semantics are similar to Pytorch, functional API is similar toJAX, tensor operators code is heavily derived fromtinygrad.Example:importslopedeff(x):y=x*2.0returny.sum()x=slope.tensor([1.,2.,3.])gf_x=slope.grad(f)(x)print(f"{gf_x=}")gf_x=<Tensor: val=
[2. 2. 2.]
shape=(3,), dtype=float32, device='cpu:0'>Installgit clone https://github.com/radenmuaz/slope
cd slope
pip install -e .Or you can just copysrc/slopeto your projects.FeaturesForward-mode, reverse-mode, and higher-order AD.Just-in-time compilation, with interchangeable backends, supporting CPU, CUDA and Metal:ONNX Runtime(ONNX graph); this is the default backendIREE(StableHLO MLIR)NumPy (Python code)Training and inference, examples:MLP on MNISTResNet on CIFAR-10Export jitted functionSmall (?)<3000 lines of core codeslope/core.py, after run withblack src --line-length 140Operators and procedures system33 core operators defined inslope/operators.pyUnary:exp log sin sqrt invert cast stop_gradientBinary:add mul sub div pow equal less greater maximumReduce:sum maxShape:reshape expand permute slice pad flip catInit:full arange random_normal random_uniformGeneralReduce:matmul conv gather_nd scatter_ndComposite operators system with "procedures"slope/procedures.pyProcedures are functions containing calls to operators, exposed withTensor.procedure_name(*args)syntax.Useful for definitions like:x.cos(), wheredef cos(x): return (math.pi/2 - x).sin()x.conv_transpose(w): wheredef conv_transpose(x, w): ...is a very long function.An operator can be directly implemented as code translation to backend, orfallbackto a procedure, e.g. there isconvprocedure in case the backend has no implementation for it.ExtensibleAdd new backend by defining implementation translationsslope/backendsNN moduleslope/nn.pyQuickstartThere are many examples inexamples/folder.We start by running MNIST classifier training,examples/nn/mnist_mlp.pypythonexamples/nn/mnist_mlp.pyStartingtraining...
...
Trainepoch:2,batch:299/300,loss:12.51:100%|██████████████████████████████████████████████████████████████████████|300/300[00:02<00:00,139.36it/s]Epoch2in2.15sec
Testsetaccuracy0.97By setting theSLOPE_BACKENDflag, we change the backend to eitheriree(default),onnxruntimeandnumpy.
We can also setLOG_JIT=1to verbose print the backend output.LOG_JIT=1SLOPE_BACKEND=onnxruntimepythonexamples/nn/mnist_mlp.py...

----train_stepcodegen:

<ir_version:7,opset_import:["":18,"slope":1]>
main(float[100,784]x0,float[10,100]x1,float[200,28,28]x2,float[200,10]x3,int32[]x4,float[100,784]x5,float[10,100]x6,float[]x7)=>(float[]y0,float[100,784]y2,float[10,100]y4,int32[]y5,float[100,784]y1,float[10,100]y3,float[]x7){z0_shape=Constant<value=int64[2]{200,784}>()z0=Reshape(x2,z0_shape)...y4=Sub(x1,z164)z165_fill_value=Constant<value=int32[1]{1}>()z165_squeeze_dim=Constant<value=int64[1]{0}>()z165=Squeeze(z165_fill_value,z165_squeeze_dim)y5=Add(x4,z165)}...===============Trainepoch:0,batch:58/300,loss:71.23:20%|██████████████▎|59/300[00:01<00:04,55.45it/s]Environment flagsput this before the command to set# prints the jitted code
LOG_JIT=1 

# set device
DEFAULT_DEVICE=cpu
DEFAULT_DEVICE=cuda
DEFAULT_DEVICE=metal

# set backend
SLOPE_BACKEND=iree # iree backend (default)
SLOPE_BACKEND=onnxruntime # onnxruntime backend
SLOPE_BACKEND=numpy # numpy backend (extremely SLOW)Slope internals tutorialSlope has familiar Pytorch-like syntaxMost of the things familiar in Pytorch works in Slope, probably.importslopex=slope.ones(2,5)w=slope.arange(15,dtype=slope.float32).reshape(5,3)b=slope.tensor([1.,2.,3.],dtype=slope.float32)y=x@w+bprint(y)Every operations are compiled with slope.jitOperation calls are jitted as individual programs eagerly.
Try running this on terminal:LOG_JIT=1python-c"import slope; print(slope.ones(3)*2)"...
----full_shape__lp__rp__fill_value_2_dt_0_dtype_<DType:float32>_device_<Device:'cpu:0'>_codegen:

func.func@main()->(tensor<f32>){%y0="stablehlo.constant"(){value=dense<2.0>:tensor<f32>}:()->(tensor<f32>)"func.return"(%y0):(tensor<f32>)->()}...# plenty of outputs....To prevent eager jit, write code in function and use slope.jit.
Then call the function@slope.jitdeff(x):y=x*xreturny# Alternative way to jit# f = slope.jit(f)x=slope.full((1,),2.)y=f(x)print(y)To see the actual code:jit_object=f.lower(x)# slope Program intermediate representationprint(jit_object.program)# backend codeprint(jit_object.code)deff(x0):# [1, f32] -> [1, f32]y0=slope.mul(x0,x0)# ([1, f32], [1, f32]) -> [1, f32]returny0func.func@main(%x0:tensor<1xf32>)->(tensor<1xf32>){%y0="stablehlo.multiply"(%x0,%x0):(tensor<1xf32>,tensor<1xf32>)->(tensor<1xf32>)"func.return"(%y0):(tensor<1xf32>)->()}Derivatives and gradientsSlope has several AD functions, likeslope.jvpslope.vjpandslope.gradTo do the usual backprop things:deff(x,w):y=x@wreturnydefloss_fn(x,w,y):y_hat=f(x,w)return((y_hat-y)**2).sum()gloss_fn=slope.value_and_grad(loss_fn,argnums=(1,))@slope.jitdeftrain_step(x,w,y,lr):loss,gw=gloss_fn(x,w,y)w=w-lr*gwreturnloss,wN=50x=slope.randn(N,2)y=slope.randn(N,1)w=slope.randn(2,1)lr=slope.tensor([0.001])foriinrange(10):loss,w=train_step(x,w,y,lr)print(i,loss.numpy())0 102.412125
1 88.60157
2 78.322815
3 70.644066
4 64.883766
5 60.54294
6 57.25553
7 54.75257
8 52.83598
9 51.359528ContributingFork this repo and hack, and maybe do a PR, too many things need to be done (see Roadmap)
Idk everything is flaky and I am still experimenting and doing many API changes, maybe later I will open a new github repo.RoadmapDocsSymbolic shape inferenceDynamic shape jitOptimizer filter frozen paramsvmap vjp and jvp to compute jacobian and hessianiree backend currently has fixed seed random, implement threefry and JAX-like randommake things fastllama (gpt) trainingwhisper inferencecore tests, operators tests on all Trace types