nn.Parameter equivalent in mlx

[henrhie]

How do you achieve nn.Parameter (which is used to create a parameter for a module in PyTorch) in mlx

[angeloskath]

Any public member of a nn.Module (ie its name not starting with an underscore) that is an array is a trainable parameter.

You can have "frozen" parameters which will not be returned by trainable_parameters() by calling freeze. Or you can have what is effectively constants by setting private member variables starting with an underscore. For instance:

import mlx.core as mx
import mlx.nn as nn

class Foo(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = mx.array(0.)
        self.b = mx.array(1.)
        self._c = mx.array(2.)

f = Foo()

print(f.trainable_parameters())
# {'a': array(0, dtype=float32), 'b': array(1, dtype=float32)}
print(f.parameters())
# {'a': array(0, dtype=float32), 'b': array(1, dtype=float32)}

f.freeze(keys=["b"])
print(f.trainable_parameters())
# {'a': array(0, dtype=float32)}
print(f.parameters())
# {'a': array(0, dtype=float32), 'b': array(1, dtype=float32)}

print(f._c)
array(2, dtype=float32)

[henrhie]

makes sense. Thanks!