Is it possible to generate Python from code quotations?

[pkese]

Here's an actual real world problem that Fable.Python can solve...

PyTorch has a JIT mode, that can make neural network models run much faster: Speeding up model training with PyTorch JIT

The way JIT works is:
Instead of interpreting Python neural network code and pushing tensor ops to the GPU line by line, you produce TorchScript code containing neural network model's AST.
TorchScript is restricted subset of Python that is simple enough, that it can be parsed by some C++ code inside Torch and then compiled into a CUDA kernel. And then that CUDA kernel code gets executed in a single call to the GPU making everything much faster (hence JIT).

Problem is that TorchScript (in Python world) is pain to produce. There are two ways:

1. write restricted subset of Python by hand (error prone, because you need to figure out which subset of Python is allowed),
2. write normal Python code and then run that code through a runtime tracer to reconstruct the original model AST from the trace. The problem is that tracing won't explore all if/then/else branches unless inputs are carefully designed. It also needs to be told which tensor dimensions are fixed and which are variable, etc.

F# solves this: unlike Python, F# makes it easy to produce ASTs without tracing.

So the idea is to use F# code quotations / computational expressions / reflected definitions of the F# AST code, then convert the AST into the TorchScript subset of Python, load that it into Torch to render fast CUDA kernel and then call that kernel directly.

Most of the infrastructure is already present: TorchSharp project contains the dotnet bindings to C++ Torch libraries (the same stuff that PyTorch and Lua Torch are built upon).
C++ Torch can load .pt files containing TorchScript models.
TorchSharp could easily load .pt files as well (the API is there), but unfortunately the C# code is lacking the Python parser and unpickler to extract Python method signatures from TorchScript, so that part is not supported yet. But if we generated the Python code ourselves from F#, we wouldn't need to parse that Python anyway, because we knew the method signatures all along from F# AST.

So the question is, which parts of Fable.Python can be reused for this?
At the minimum, I assume that Python AST parts can come handy, but I hope one can start somewhere closer to F# AST and reuse more of Fable infrastructure.

I'm just assuming that code quotations would be preferred here, because the tensor model hot-spot code that needs to get compiled into CUDA is usually just a few lines of code compared to whole project, that includes data loading, training loop, setting up the optimizers, etc... i.e. stuff that runs just fine in F# TorchSharp and can't be put on the GPU anyway. We could potentially parse the whole source file using FCS and throw everything else but the model code away.

Any thoughts?

[pkese]

One more thing:
this is useful even without full Python standard library being implemented.

Tensor code is mostly calling just:

• mathematical operators
• methods on tensors, indexing and slicing
• optionally pushing tensors to a list (in a loop) before stacking them up into a result tensor.

See TorchScript Language Reference

[dbrattli]

@pkese This looks interesting, but I'm not sure about the state of quotation support in Fable. We will need to investigate fable-compiler/Fable#1839

[pkese]

I don't think that full quotation support (as in in Fable runtime) would be necessary, because the hosting app is just normal dotnet CLR with all F# core libraries being accessible. It's just quotations that need to be converted into something that could be pushed into Fable compiler.

I'm experimenting with mapping Quotations.Expr into Fable.AST:

let logPoisson =
    <@
        fun (k:torch.Tensor) (mu:torch.Tensor) ->
            let bcast = torch.broadcast_tensors(k, mu)
            let k = bcast[0]
            let mu = bcast[1]
            let logPmf = k.xlogy(mu) - (k + torch.tensor 1).lgamma() - mu
            -logPmf
    @>

... and manage to Fable-compile the above into Python code below (I haven't implemented Fable's DeclaredTypes yet, so generated types are all Any):

def arrow_1(k: Any=None) -> Callable[[Any], Any]:
    def arrow_0(mu: Any=None) -> Any:
        bcast : Any = torch.broadcast_tensors([k, mu])
        k : Any = bcast[0]
        mu : Any = bcast[1]
        logPmf : Any = (k.xlogy(mu) - k + torch.tensor(1, None, False).lgamma()) - mu
        return -logPmf

    return arrow_0

However I'm not sure this is the best approach. It seems some AST conversion gets done in FSharp2Fable code, which is operating on FSharpExpr (typed AST), i.e. a step before I plug my AST into Fable.
(Slightly more complex examples don't compile correctly, e.g. F# functions with multiple comma-separated parameters get encoded into python as single tuple parameter).

Problem is that typed FSharpExpr AST is much more detailed than Quotations.Expr which makes it hard to convert in that direction.

FSharp.Quotations.Compiler made a work-around by converting Quotations.Expr into untyped SynExpr AST (https://github.com/eiriktsarpalis/QuotationCompiler/blob/master/src/QuotationCompiler/Compiler.fs) and feeds that untyped AST into Compiler Service to get compiled IL (https://github.com/eiriktsarpalis/QuotationCompiler/blob/master/src/QuotationCompiler/QuotationCompiler.fs#L131).
But the way Compiler Service is structured, one can get from untyped AST to IL, but there's no way to get the intermediate typed AST.

[pkese]

Btw, I've noticed that Python generated parentheses are wrong:

logPmf : Any = (k.xlogy(mu) - k + torch.tensor(1, None, False).lgamma()) - mu

should be

logPmf : Any = k.xlogy(mu) - (k + torch.tensor(1, None, False)).lgamma() - mu

I can probably fix that on my side with using witnesses and generating proper tensor.add(a,b) instead of a + b, but it might still be interesting to understand why this is happening on Python side. Would me adding correct types instead of Any help?

[dbrattli]

What did the original F# code for this line look like? Or could you make a minimal repro without the deps?

[dbrattli]

FYI: @alfonsogarciacaro for info about quotation handling. He might provide more input here.

[pkese]

@dbrattli here's a repro:

let s = "Py_" + ("str" + "ing").ToUpper() + ";"

compiles into

s : str = ("Py_" + "str" + "ing".ToUpper()) + ";"

... and it looks like I probably have something wrong with how I'm calling compiler, because F#'s string .ToUpper() doesn't get replaced with Python's .upper().

But otherwise, yeah, I'd appreciate any hints or directions from folks that have been here before me.
It's my first encounter with FCS and Fable and there's lots of internals I'm not aware of, so I might be digging in the wrong direction.

[dbrattli]

@pkese What version of Fable are you running? dotnet fable-py --version

When I try the example I get:

s : str = ("Py_" + "string".upper()) + ";"

❯ dotnet fable-py --version
4.0.0-alpha-032

Try: dotnet dotnet tool update fable-py --version 4.0.0-alpha-032

[pkese]

I'm using the beyond branch of Fable and I'm creating the compiler from my own code.

There's no official way (or at least I didn't find it) to make Fable start compilation from some provided AST ... Fable assumes there's a project with .fs files and wants to parse that project and read those files.

So I initialized CompilerImpl by hand and I created Fable.File and put the AST into that File and feed that to Fable.Transforms.Fable2Python.Compiler.transformFile.

[dbrattli]

@alfonsogarciacaro Do you have any ideas about this?

[pkese]

So for the context (after spending a few days on this), I'm trying to evaluate my options.

Currently I'm:

• starting with Quotations.Expr
• manually translating that into Fable.AST
• feeding that AST to Fable.Transforms.Fable2Python.Compiler

This appears to be working for simple cases and it gets me quite far - CUDA kernel code probably won't need full F# language support anyway. However I'm skipping FSharp2Fable step which I assume I'll have to partially re-implement manually (not being familiar enough with Fable internals, I don't know which transforms get applied on which step, so these are just my assumptions).

A thing that I haven't yet figured out is also how to plug in replacements. In few places, I'll have to translate between F#'s TorchSharp library calls and their PyTorch counterparts (differences are minimal).
Issue with replacements is that currently I'm not even being able to get F#'s string.ToUpper() to map to Python's string.upper() even though this replacement is being defined and applied on Fable.AST. I probably haven't configured everything properly, or I'm skipping a step in Fable (which?)... I'm assuming that probably that part of Fable.AST translation happens before Fable2Python compiler.

---

The other option would be to replicate what FSharp.Quotations.Compiler is doing:

• start with Quotations.Expr
• convert that into untyped SynExpr (this code already exists in FSharp.Quotations.Compiler)
• feed SynExpr into FCS.Checker to get typed FSharpExpr AST
• feed typed AST into Fable
• profit from having full F# language support without having to manually re-implement FSharp2Fable.

The problem here is that this would need a fork of FCS, because FCS doesn't expose any APIs that would digest untyped SynExpr AST and spit out typed FSharpExpr AST.

---

The third option would be to compile the whole project source code with Fable and then cherry-pick and extract just bits and pieces that need to be passed to CUDA. The issue with that is that this would probably make it difficult to work with dynamic Jupyter Notebook environments, which - for many - is the standard way to develop AI code.

---

Also my issue is that I'm just getting familiar with the whole FCS / Fable compiler infrastructure - and the learning curve is quite steep. Especially because I'm trying to plug my AST into the Fable workflow at locations, where it probably wasn't intended (and I'm wondering if any of this can maybe be made more approachable with Fable 4?).
But the positive side is that I'm solving my own pain and I have some motivation to work on this - i.e. it's not an abstract thing that no-one would care about.

I'll appreciate any hints or directions.

[alfonsogarciacaro]

Sorry for the late reply. I need to check to understand the problem in full but from a quick read this looks like a good use case for Fable plugins. As mentioned above quotations are not supported by Fable yet, the problem is not quotation itself but in order to make them useful you need almost "full" reflection support which prevents tree shaking and produces much bigger bundle sizes. This is way #1839 is not merged yet.

But quotations are only necessary if you need to manipulate the AST in runtime. If everything can be resolved at compile time a plugin should be enough. The plugin receives the AST for the decorated method (it can also inspect the AST of the full file) and/or calls to that method, together with a helper from the compiler and returns the transformed AST. If it's enough for you to do the transformations in the Fable AST, Fable2Python can later transform that to Python. But if you need to generate your custom code you can just build a string an return Emit. You can check Feliz ReactComponent plugin source code for reference.

The problem with plugins in beyond is we're still changing the AST, so plugins are not guaranteed to work until we release a more stable version (though that's the same issue if you fork the compiler). But hopefully they will only require minor changes.