Clean up axis hooks in nnx.Variable

[IvyZX]

• Swapped the order of axis index and name in hooks. This is because in pure JAX, an add/remove axis always needs an index, but perhaps doesn't need a name.

  • This will allow hooks to be triggered on every lifted transform, even if sharding annotations is not part of the model.

• Simplified the boilderplate code of creating hooks in nnx.Variable.

• Add a simple test on add/remove axis.

[cgarciae]

The PR looks good, swapping the arguments makes sense.

We don't have to do this here, but I was wondering if we could conditionally add the hooks if at least one hook is present to avoid adding useless metadata e.g:

if get_value_hooks:
  self.get_value_hooks = get_value_hooks

and check for existence on usage e.g.

if hasattr(self, 'get_value_hooks'):
  for hook in self.get_value_hooks:
    value = hook(self, value)

[IvyZX]

The PR looks good, swapping the arguments makes sense.

We don't have to do this here, but I was wondering if we could conditionally add the hooks if at least one hook is present to avoid adding useless metadata e.g:

if get_value_hooks:
  self.get_value_hooks = get_value_hooks

and check for existence on usage e.g.

if hasattr(self, 'get_value_hooks'):
  for hook in self.get_value_hooks:
    value = hook(self, value)

You know what, I tried this, but it ended up with infinite recursion because hasattr(self, 'get_value_hooks') calls __getattr__ which calls hasattr(self, 'get_value_hooks') again... We are Pythonic indeed

We can make get_value_hooks and set_value_hooks always available, but that creates diff between them and other hooks so I'm not so sure

[cgarciae]

Ah yes I've faced this, we need to we should fix it by raising an AttributeError manually inside __getattr__.