Cannot return map from Axon.nx

[jonatanklosko]

model =
  Axon.input("x")
  |> Axon.nx(fn x -> %{x: x} end)

Axon.predict(model, %{}, %{"x" => Nx.tensor(1)})

Fails with

** (Axon.CompileError) exception found when compiling layer anonymous fn/2 in Axon.nx/3 named nx_0:

    ** (ArgumentError) unable to build tensor expression, expected a tensor or a number, got: %{x: #Nx.Tensor<
        f32
        
        Nx.Defn.Expr
        parameter a:0   s64
        b = as_type a   f32
      >}
        (nx 0.7.0-dev) lib/nx/defn/expr.ex:1333: Nx.Defn.Expr.to_expr/1
    

(pass debug: true to build/compile see where the layer was defined)


Compiling of the model was initiated at:

    (nx 0.7.0-dev) lib/nx/defn/compiler.ex:158: Nx.Defn.Compiler.runtime_fun/3
    (nx 0.7.0-dev) lib/nx/defn/evaluator.ex:82: Nx.Defn.Evaluator.precompile/3
    (nx 0.7.0-dev) lib/nx/defn/evaluator.ex:60: Nx.Defn.Evaluator.__compile__/4
    (nx 0.7.0-dev) lib/nx/defn/evaluator.ex:53: Nx.Defn.Evaluator.__jit__/5
    (nx 0.7.0-dev) lib/nx/defn.ex:433: Nx.Defn.do_jit_apply/3

Regression from 235a030.