Prettify the printing

.typos.toml

@@ -0,0 +1,2 @@
+[default.extend-words]
+numer = "numer"

examples/GravitationalWaveForm/main.jl

@@ -163,7 +163,7 @@ function compute_waveform(dt::T, soln, mass_ratio, model_params=nothing) where {
  m₁ = mass_ratio * m₂
 
  orbit₁, orbit₂ = one2two(orbit, m₁, m₂)
- waveform = h_22_strain_two_body(dt, orbit1, mass1, orbit2, mass2)
+ waveform = h_22_strain_two_body(dt, orbit₁, m₁, orbit₂, m₂)
  else
  waveform = h_22_strain_one_body(dt, orbit)
  end
@@ -183,11 +183,11 @@ end
 function RelativisticOrbitModel(u, (p, M, e), t)
  χ, ϕ = u
 
- number = (p - 2 - 2 * e * cos(χ)) * (1 + e * cos(χ))^2
+ numer = (p - 2 - 2 * e * cos(χ)) * (1 + e * cos(χ))^2
  denom = sqrt((p - 2)^2 - 4 * e^2)
 
- χ̇ = number * sqrt(p - 6 - 2 * e * cos(χ)) / (M * (p^2) * denom)
- ϕ̇ = number / (M * (p^(3 / 2)) * denom)
+ χ̇ = numer * sqrt(p - 6 - 2 * e * cos(χ)) / (M * (p^2) * denom)
+ ϕ̇ = numer / (M * (p^(3 / 2)) * denom)
 
  return [χ̇, ϕ̇]
 end
@@ -260,11 +260,11 @@ function ODE_model(u, nn_params, t)
  ## it, however, in general, we should use `st` to store the state of the neural network.
  y = 1 .+ nn_model([first(u)], nn_params)
 
- number = (1 + e * cos(χ))^2
+ numer = (1 + e * cos(χ))^2
  denom = M * (p^(3 / 2))
 
- χ̇ = (number / denom) * y[1]
- ϕ̇ = (number / denom) * y[2]
+ χ̇ = (numer / denom) * y[1]
+ ϕ̇ = (numer / denom) * y[2]
 
  return [χ̇, ϕ̇]
 end

examples/NeuralODE/main.jl

@@ -236,11 +236,11 @@ x = gpu_device()(ones(Float32, 28, 28, 1, 3));
 
 @code_warntype model_stateful(x, ps_stateful, st_stateful)
 
+# Note, that we still recommend using this layer internally and not exposing this as the
+# default API to the users.
+
 # Finally checking the compact model
 
 model_compact, ps_compact, st_compact = create_model(NeuralODECompact)
 
 @code_warntype model_compact(x, ps_compact, st_compact)
-
-# Note, that we still recommend using this layer internally and not exposing this as the
-# default API to the users.

examples/SimpleRNN/main.jl

@@ -108,8 +108,7 @@ end
 function SpiralClassifierCompact(in_dims, hidden_dims, out_dims)
  lstm_cell = LSTMCell(in_dims => hidden_dims)
  classifier = Dense(hidden_dims => out_dims, sigmoid)
- return @compact(; lstm_cell=lstm_cell,
- classifier=classifier) do x::AbstractArray{T, 3} where {T}
+ return @compact(; lstm_cell, classifier) do x::AbstractArray{T, 3} where {T}
  x_init, x_rest = Iterators.peel(Lux._eachslice(x, Val(2)))
  y, carry = lstm_cell(x_init)
  for x in x_rest

src/helpers/compact.jl

@@ -200,8 +200,6 @@ function __compact_macro_impl(_exs...)
  vars = map(first ∘ Base.Fix2(getproperty, :args), kwexs)
  fex = supportself(fex, vars, splatted_kwargs)
 
- display(fex)
-
  # assemble
  return esc(:($CompactLuxLayer{$dispatch}($fex, $name, ($layer, $input, $block),
  (($(Meta.quot.(splatted_kwargs)...),), ($(splatted_kwargs...),)); $(kwexs...))))
@@ -346,9 +344,8 @@ function CompactLuxLayer{dispatch}(
  if is_lux_layer
  setup_strings = merge(setup_strings, NamedTuple((name => val,)))
  else
- setup_strings = merge(setup_strings,
- NamedTuple((name => sprint(
- show, val; context=(:compact => true, :limit => true)),)))
+ setup_strings = merge(
+ setup_strings, NamedTuple((name => __kwarg_descriptor(val),)))
  end
  end
 
@@ -410,3 +407,20 @@ function Lux._big_show(io::IO, obj::CompactLuxLayer, indent::Int=0, name=nothing
  end
  return
 end
+
+function __kwarg_descriptor(val)
+ val isa Number && return string(val)
+ val isa AbstractArray && return sprint(Base.array_summary, val, axes(val))
+ val isa Tuple && return "(" * join(map(__kwarg_descriptor, val), ", ") * ")"
+ if val isa NamedTuple
+ fields = fieldnames(typeof(val))
+ strs = []
+ for fname in fields[1:min(length(fields), 3)]
+ internal_val = getfield(val, fname)
+ push!(strs, "$fname = $(__kwarg_descriptor(internal_val))")
+ end
+ return "@NamedTuple{$(join(strs, ", "))" * (length(fields) > 3 ? ", ..." : "") * "}"
+ end
+ val isa Function && return sprint(show, val; context=(:compact => true, :limit => true))
+ return lazy"$(nameof(typeof(val)))(...)"
+end

test/helpers/compact_tests.jl

@@ -180,7 +180,7 @@
  return w(x .* s)
  end
  expected_string = """@compact(
- x = randn(32),
+ x = 32-element Vector{Float64},
  w = Dense(32 => 32), # 1_056 parameters
  ) do s 
  return w(x .* s)
@@ -197,8 +197,8 @@
  end
  expected_string = """@compact(
  w1 = Model(32)(), # 1_024 parameters
- w2 = randn(32, 32),
- w3 = randn(32),
+ w2 = 32×32 Matrix{Float64},
+ w3 = 32-element Vector{Float64},
  ) do x 
  return w2 * w1(x)
  end # Total: 2_080 parameters,