feat: nanoGPT implementation using Reactant

[avik-pal]

Needs

• NNlib.gather support in Reactant: feat: partial NNlib.gather support + better indexing support EnzymeAD/Reactant.jl#252
• NNlib.make_causal_mask: feat: add support for NNlib make causal mask EnzymeAD/Reactant.jl#249
• adjoint for dynamic_gather xref Jmp/scatter gather EnzymeAD/Enzyme-JAX#142

error: could not compute the adjoint for this operation %51 = "stablehlo.dynamic_gather"(%15, %0, %12) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1], index_vector_dim = 1>}> : (tensor<64x64xf32>, tensor<64xi64>, tensor<2xi64>) -> tensor<64x64xf32>
error: could not compute the adjoint for this operation %47 = "stablehlo.dynamic_gather"(%14, %46, %12) <{dimension_numbers = #stablehlo.gather<offset_dims = [0], collapsed_slice_dims = [1], start_index_map = [1], index_vector_dim = 1>}> : (tensor<64x67xf32>, tensor<8192xi64>, tensor<2xi64>) -> tensor<64x8192xf32>
ERROR: LoadError: "failed to run pass manager on module"

[Vaibhavdixit02]

Suggested change

	) |> dev
	) .|> dev

[avik-pal]

why is this needed? broadcasting over the dataloader will avoid constructing the DeviceIterator

[Vaibhavdixit02]

Atleast, on GPU without this the labels weren't being transferred to the device

[Vaibhavdixit02]

Suggested change

	Array(first(model_compiled(testX, ps, Lux.testmode(st)))), testY
	first(model_compiled(testX, ps, Lux.testmode(st))), testY

[avik-pal]

We need to compile the loss_fn for the non-Array version to work

[Vaibhavdixit02]

Ah yeah this might not be needed for your case here

Copilot reviewed 3 out of 4 changed files in this pull request and generated no suggestions.

Files not reviewed (1)

• examples/NanoGPT/main.jl: Language not supported