ModularM
Modular17mo ago
5 replies
Hylke

How to define nested dict function

(This is a shortened cross post of my stackoverflow question).

How do I define a function in Mojo where the argument or return value is a nested structure, like a dictionary. But where the exact structure is not known at compile time because, for example, it is read from disk.

For example, how do I type annotate params in the function neural_network to represent a nested dict to keep the model parameters (i.e., a pytree).
network_params = {
  'linear1': {'weights': Tensor(...), 'bias': 1.0},
  ...
}

fn neural_network(params: Dict, x: Tensor) -> Float64:
  x = params['linear1']['weights'] @ x + params['linear1']['bias']
  ...
  return x

x_input = Tensor(...)
y = neural_network(network_params, x_input)

Thanks so much!
Stack Overflow
How do I define a function in Mojo where the argument is a nested structure, like a dictionary. But where the exact structure is not known at compile time because, for example, it is read from disk...
How to define a nested dictionary argument of arbitrary depth
Was this page helpful?