function Torchx.__jit__/5 is undefined or private

[hiramegl]

Hi,

I wonder if this error is something that needs to be fixed in Nx or in Torchx:

** (UndefinedFunctionError) function Torchx.__jit__/5 is undefined or private
    (torchx 0.7.1) Torchx.__jit__(#Function<134.37423472/2 in Nx.Defn.Compiler.fun/2>, [{#Nx.Tensor<
     f32[128][52]

     Nx.Defn.Expr
     parameter a:0   f32[128][52]
   >, #Nx.Tensor<
     f32[128][1]

     Nx.Defn.Expr
     parameter a:1   f32[128][1]
   >}, %{}], #Function<147.37423472/1 in Nx.Defn.Compiler.prepare_options/2>, [[#Function<0.82037917/0 in Nx.LazyContainer.Nx.Tensor.traverse/3>, #Function<0.82037917/0 in Nx.LazyContainer.Nx.Tensor.traverse/3>]], [])
    (nx 0.7.1) lib/nx/defn.ex:452: Nx.Defn.do_jit_apply/3
    (axon 0.6.1) lib/axon/loop.ex:1780: Axon.Loop.init_loop_state/8
    (stdlib 5.1.1) timer.erl:270: :timer.tc/2
    (axon 0.6.1) lib/axon/loop.ex:1646: Axon.Loop.run/4

Thanks for your help in advance!

Best regards,

/Hiram

[josevalim]

It seems you are configuring Torchx as a compiler. Torchx is not a compiler, only a backend. If you can share the code that triggers the error, we can provide more precise feedback.

[hiramegl]

Hi Jose,

Yes, that's right! I was using Torchx as compiler. I thought I needed to use it when I switched from EXLA backend to Torchx backend. Now the loop runs without problem.

model
|> Axon.Loop.trainer(:binary_cross_entropy, :adam)
|> Axon.Loop.run(
  training_data,
  %{},
  compiler: EXLA, # Torchx,
  epochs: 10)

Thanks a lot Jose!

Best regards,

/Hiram