Reactant compatibility

[AntonOresten]

There should be no breaking changes. Layer types have been given type parameters for type stability and compilability. Layer types are also now exported.

Compiling with Reactant

TransformerBlock

This seems to work fine

using Jjama3, Reactant, MLDataDevices
Reactant.set_default_backend("cpu")
const xdev = reactant_device()
n = 10
tokens = rand(1:256, n, 1) |> xdev;
model = Transformer(256, 1024, 4, 8, 8, 256, 2048) |> xdev;
block = model.layers[1] |> xdev;
h = @jit model.tok_embeddings(tokens);
@jit block(h, 0, model.rope[1:n], false);

Transformer

This fails for some reason:

julia> @jit model(tokens, 0);
ERROR: BoundsError: attempt to access Ptr{Nothing} at index [1]
Stacktrace:
 [1] traced_getfield
   @ ~/.julia/packages/Reactant/m1CaM/src/Compiler.jl:17 [inlined]
 [2] macro expansion
   @ ~/.julia/packages/Reactant/m1CaM/src/Compiler.jl:799 [inlined]
 [3] (::Reactant.Compiler.Thunk{…})(::ConcreteRArray{…}, ::Int64)
   @ Reactant.Compiler ~/.julia/packages/Reactant/m1CaM/src/Compiler.jl:815
 [4] top-level scope
   @ ~/.julia/packages/Reactant/m1CaM/src/Compiler.jl:447
Some type information was truncated. Use `show(err)` to see complete types.

Buuut.... it's not an implementation thing:

julia> (::Transformer)() = nothing

julia> @jit model()
ERROR: BoundsError: attempt to access Ptr{Nothing} at index [1]
Stacktrace:
 [1] traced_getfield
   @ ~/.julia/packages/Reactant/m1CaM/src/Compiler.jl:17 [inlined]
 [2] macro expansion
   @ ~/.julia/packages/Reactant/m1CaM/src/Compiler.jl:799 [inlined]
 [3] (::Reactant.Compiler.Thunk{…})()
   @ Reactant.Compiler ~/.julia/packages/Reactant/m1CaM/src/Compiler.jl:815
 [4] top-level scope
   @ ~/.julia/packages/Reactant/m1CaM/src/Compiler.jl:447
Some type information was truncated. Use `show(err)` to see complete types.

So probably has to do with the fields or type parameters? My bet is the transformer blocks, which is currently a AbstractVector{<:TransformerBlock}. We might opt for using an NTuple or Chain.

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 106 lines in your changes missing coverage. Please review.

Project coverage is 0.00%. Comparing base (0b50d41) to head (aecca48).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/layers.jl	0.00%	54 Missing ⚠️
src/cache.jl	0.00%	20 Missing ⚠️
src/model.jl	0.00%	14 Missing ⚠️
src/sampling.jl	0.00%	10 Missing ⚠️
src/utils.jl	0.00%	8 Missing ⚠️

Additional details and impacted files

@@          Coverage Diff          @@
##            main     #15   +/-   ##
=====================================
  Coverage   0.00%   0.00%           
=====================================
  Files          5       6    +1     
  Lines        299     288   -11     
=====================================
+ Misses       299     288   -11     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[murrellb]

Nice. Can you split out a "General type stability/code improvements" PR that we can merge right away (maybe after some timing tests), and try the reactant stuff on top of that?

[AntonOresten]

This PR ended up being more of a general refactor, moving away from the style of original Python implementations to a more type-stable Julian setup, which should be compatible with Reactant, but that is not stable yet.