[nnx] Add specific model typing for nnx.Optimizer

[marcelroed]

What does this PR do?

This PR adds a type variable for nnx.Optimizer.model to make type inference and IDE autocomplete better.

Before:

class MyModel(nnx.Module):
    ...

optimizer = nnx.Optimizer(MyModel(), tx=...)

optimizer.model  # inferred type: nnx.Module

def use_optimizer(optimizer:​ nnx.Optimizer):
    optimizer.model  # inferred type: nnx.Module

After PR:

class MyModel(nnx.Module):
    ...

optimizer = nnx.Optimizer(MyModel(), tx=...)

optimizer.model  # inferred type: MyModel

def use_optimizer(optimizer:​ nnx.Optimizer[MyModel]):
    optimizer.model  # inferred type: MyModel

# Can also be used as before with generic type hints
def use_optimizer(optimizer: nnx.Optimizer):
    optimizer.model  # inferred type: nnx.Module

Checklist

• This PR fixes a minor issue.

[cgarciae]

@marcelroed thanks for doing this!